diff --git a/api/Makefile b/api/Makefile index 641b5afe..2bcfaf55 100644 --- a/api/Makefile +++ b/api/Makefile @@ -16,9 +16,13 @@ vet: fmt .PHONY: vet build: vet - go build + go build -v ./... .PHONY: build +test: + go test -v ./... +.PHONY: test + install: go install github.com/air-verse/air@latest go mod tidy diff --git a/api/artifact_instruction.txt b/api/artifact_instruction.txt deleted file mode 100644 index 0c4d05ab..00000000 --- a/api/artifact_instruction.txt +++ /dev/null @@ -1,147 +0,0 @@ -ARTIFACT CREATION GUIDELINES - MANDATORY COMPLIANCE REQUIRED: - -⚠️ CRITICAL: These formatting rules are REQUIRED for artifact rendering. Deviation will cause display failures. - -## MANDATORY ARTIFACT FORMATS (EXACT SYNTAX REQUIRED): - -### HTML Applications: -```html -[Complete HTML content with inline CSS and JavaScript(Preact)] -``` - -### SVG Graphics: -```svg -[Complete SVG markup] -``` - -### Mermaid Diagrams: -```mermaid -[Mermaid diagram syntax] -``` - -### JSON Data: -```json -[Valid JSON data] -``` - -### Executable Code: -```javascript -[JavaScript/TypeScript code] -``` - -```python -[Python code] -``` - -## FORMATTING COMPLIANCE CHECKLIST: - -✅ Comment MUST be on the SAME LINE as opening ``` -✅ Use EXACT format: `` or `` -✅ Include descriptive, specific title explaining functionality -✅ No extra spaces or characters in comment syntax -✅ Complete, self-contained code within blocks - -❌ COMMON ERRORS TO AVOID: -- Comment on separate line from ``` -- Missing or incorrect comment format -- Generic titles like "Code" or "Example" -- Incomplete or broken code -- External dependencies in HTML artifacts - -## WHEN TO CREATE ARTIFACTS: - -### ALWAYS create artifacts for: -- Interactive web applications, forms, games, tools -- Data visualizations, charts, graphs, dashboards -- Diagrams, flowcharts, visual representations -- Working code examples demonstrating functionality -- Calculators, converters, utility applications -- Rich data displays or formatted outputs -- Any content meant to be rendered/executed - -### NEVER create artifacts for: -- Simple text responses -- Code snippets for reference only -- Incomplete or pseudo-code -- Content requiring external files - -## HTML ARTIFACT STANDARDS: - -### REQUIRED STRUCTURE: -```html - - - - - - [App Title] - - - - - - - -``` - -### TECHNICAL REQUIREMENTS: -- Use Preact with HTM: `import { html, render } from 'https://unpkg.com/htm/preact/standalone.module.js'` -- Modern ES6+ syntax only -- Responsive design with proper viewport meta -- Semantic HTML5 elements -- Accessible UI with proper ARIA labels -- Error handling for user interactions - -## EXECUTABLE CODE STANDARDS: - -### JavaScript/TypeScript FEATURES: -- Output: console.log(), console.error(), console.warn() -- Graphics: createCanvas(width, height) for visualizations -- Libraries: `// @import libraryName` (lodash, d3, chart.js, moment, axios, rxjs, p5, three, fabric) -- Return values automatically displayed -- Built-in timeout and resource monitoring - -### Python FEATURES: -- Output: print() for all results (auto-captured) -- Plotting: matplotlib plots auto-displayed as PNG -- Libraries: numpy, pandas, matplotlib, scipy, scikit-learn, requests, seaborn, plotly -- Memory and execution monitoring included -- No file/network access (sandboxed) - -## QUALITY ASSURANCE: - -### PRE-SUBMISSION CHECKLIST: -1. ✅ Verify exact comment syntax on same line as ``` -2. ✅ Test all interactive functionality -3. ✅ Ensure complete self-contained code -4. ✅ Validate responsive design (HTML) -5. ✅ Confirm proper error handling -6. ✅ Check accessibility features -7. ✅ Verify cross-browser compatibility - -### ARTIFACT TITLE GUIDELINES: -- Be specific and descriptive -- Include primary function/purpose -- Avoid generic terms -- Examples: - - ✅ "Interactive BMI Calculator with Health Recommendations" - - ✅ "Real-time Stock Price Chart with Technical Indicators" - - ❌ "Calculator" - - ❌ "Chart" - -## RENDERER BEHAVIOR: - -The artifact viewer uses specialized renderers: -- **HTML**: Full browser environment with Preact support -- **SVG**: Native SVG rendering with interactive capabilities -- **Mermaid**: Diagram engine with theme support -- **JSON**: Formatted tree view with syntax highlighting -- **JavaScript**: Node.js-like environment with canvas support -- **Python**: Scientific computing sandbox with plot display - -⚠️ FINAL REMINDER: Artifacts that don't follow these exact formatting rules will fail to render. Always double-check syntax before submission. \ No newline at end of file diff --git a/api/dto/errors.go b/api/dto/errors.go index 301be71f..5ec82240 100644 --- a/api/dto/errors.go +++ b/api/dto/errors.go @@ -7,7 +7,7 @@ import ( "encoding/json" "errors" "fmt" - "log" + "log/slog" "net/http" "sort" "strings" @@ -225,7 +225,7 @@ var ( // Deprecated: use ErrClaudeResponseFailed instead. ErrClaudeResponseFaild = ErrClaudeResponseFailed - ErrOpenAIStreamFailed = APIError{ + ErrOpenAIStreamFailed = APIError{ HTTPCode: http.StatusInternalServerError, Code: ErrModel + "_006", Message: "Failed to stream OpenAI response", @@ -297,11 +297,11 @@ func RespondWithAPIError(w http.ResponseWriter, err APIError) { } if err.DebugInfo != "" { - log.Printf("Error [%s]: %s - %s", err.Code, err.Message, err.DebugInfo) + slog.Error("api error", "code", err.Code, "message", err.Message, "debug", err.DebugInfo) } if err := json.NewEncoder(w).Encode(response); err != nil { - log.Printf("Failed to write error response: %v", err) + slog.Info("Failed to write error response", "error", err) } } @@ -359,7 +359,7 @@ func MapDatabaseError(err error) error { } } - log.Printf("Unhandled database error: %v", err) + slog.Info("Unhandled database error", "error", err) dbErr := ErrDatabaseQuery dbErr.DebugInfo = err.Error() @@ -401,6 +401,15 @@ func WrapError(err error, detail string) APIError { apiErr.Detail = detail } } + case *APIError: + apiErr = *e + if detail != "" { + if apiErr.Detail != "" { + apiErr.Detail = fmt.Sprintf("%s: %s", detail, apiErr.Detail) + } else { + apiErr.Detail = detail + } + } default: apiErr = ErrInternalUnexpected apiErr.Detail = detail @@ -421,39 +430,39 @@ func IsErrorCode(err error, code string) bool { // --- Error catalog (for documentation) --- var ErrorCatalog = map[string]APIError{ - ErrAuthInvalidCredentials.Code: ErrAuthInvalidCredentials, - ErrAuthExpiredToken.Code: ErrAuthExpiredToken, - ErrAuthAdminRequired.Code: ErrAuthAdminRequired, - ErrAuthInvalidEmailOrPassword.Code: ErrAuthInvalidEmailOrPassword, - ErrAuthAccessDenied.Code: ErrAuthAccessDenied, - ErrResourceNotFoundGeneric.Code: ErrResourceNotFoundGeneric, - ErrResourceAlreadyExistsGeneric.Code: ErrResourceAlreadyExistsGeneric, - ErrTooManyRequests.Code: ErrTooManyRequests, - ErrChatSessionNotFound.Code: ErrChatSessionNotFound, - ErrChatFileNotFound.Code: ErrChatFileNotFound, - ErrChatModelNotFound.Code: ErrChatModelNotFound, - ErrChatMessageNotFound.Code: ErrChatMessageNotFound, + ErrAuthInvalidCredentials.Code: ErrAuthInvalidCredentials, + ErrAuthExpiredToken.Code: ErrAuthExpiredToken, + ErrAuthAdminRequired.Code: ErrAuthAdminRequired, + ErrAuthInvalidEmailOrPassword.Code: ErrAuthInvalidEmailOrPassword, + ErrAuthAccessDenied.Code: ErrAuthAccessDenied, + ErrResourceNotFoundGeneric.Code: ErrResourceNotFoundGeneric, + ErrResourceAlreadyExistsGeneric.Code: ErrResourceAlreadyExistsGeneric, + ErrTooManyRequests.Code: ErrTooManyRequests, + ErrChatSessionNotFound.Code: ErrChatSessionNotFound, + ErrChatFileNotFound.Code: ErrChatFileNotFound, + ErrChatModelNotFound.Code: ErrChatModelNotFound, + ErrChatMessageNotFound.Code: ErrChatMessageNotFound, ErrValidationInvalidInputGeneric.Code: ErrValidationInvalidInputGeneric, - ErrChatFileTooLarge.Code: ErrChatFileTooLarge, - ErrChatFileInvalidType.Code: ErrChatFileInvalidType, - ErrChatSessionInvalid.Code: ErrChatSessionInvalid, - ErrDatabaseQuery.Code: ErrDatabaseQuery, - ErrDatabaseConnection.Code: ErrDatabaseConnection, - ErrDatabaseForeignKey.Code: ErrDatabaseForeignKey, - ErrExternalTimeout.Code: ErrExternalTimeout, - ErrExternalUnavailable.Code: ErrExternalUnavailable, - ErrInternalUnexpected.Code: ErrInternalUnexpected, - ErrChatStreamFailed.Code: ErrChatStreamFailed, - ErrChatRequestFailed.Code: ErrChatRequestFailed, - ErrSystemMessageError.Code: ErrSystemMessageError, - ErrClaudeStreamFailed.Code: ErrClaudeStreamFailed, - ErrClaudeRequestFailed.Code: ErrClaudeRequestFailed, - ErrClaudeInvalidResponse.Code: ErrClaudeInvalidResponse, - ErrClaudeResponseFailed.Code: ErrClaudeResponseFailed, - ErrOpenAIStreamFailed.Code: ErrOpenAIStreamFailed, - ErrOpenAIRequestFailed.Code: ErrOpenAIRequestFailed, - ErrOpenAIInvalidResponse.Code: ErrOpenAIInvalidResponse, - ErrOpenAIConfigFailed.Code: ErrOpenAIConfigFailed, + ErrChatFileTooLarge.Code: ErrChatFileTooLarge, + ErrChatFileInvalidType.Code: ErrChatFileInvalidType, + ErrChatSessionInvalid.Code: ErrChatSessionInvalid, + ErrDatabaseQuery.Code: ErrDatabaseQuery, + ErrDatabaseConnection.Code: ErrDatabaseConnection, + ErrDatabaseForeignKey.Code: ErrDatabaseForeignKey, + ErrExternalTimeout.Code: ErrExternalTimeout, + ErrExternalUnavailable.Code: ErrExternalUnavailable, + ErrInternalUnexpected.Code: ErrInternalUnexpected, + ErrChatStreamFailed.Code: ErrChatStreamFailed, + ErrChatRequestFailed.Code: ErrChatRequestFailed, + ErrSystemMessageError.Code: ErrSystemMessageError, + ErrClaudeStreamFailed.Code: ErrClaudeStreamFailed, + ErrClaudeRequestFailed.Code: ErrClaudeRequestFailed, + ErrClaudeInvalidResponse.Code: ErrClaudeInvalidResponse, + ErrClaudeResponseFailed.Code: ErrClaudeResponseFailed, + ErrOpenAIStreamFailed.Code: ErrOpenAIStreamFailed, + ErrOpenAIRequestFailed.Code: ErrOpenAIRequestFailed, + ErrOpenAIInvalidResponse.Code: ErrOpenAIInvalidResponse, + ErrOpenAIConfigFailed.Code: ErrOpenAIConfigFailed, } // ErrorCatalogHandler serves the error catalog as JSON. diff --git a/api/handler/auth_password.go b/api/handler/auth_password.go index aa6ef921..5715678f 100644 --- a/api/handler/auth_password.go +++ b/api/handler/auth_password.go @@ -4,10 +4,10 @@ import ( "encoding/json" "net/http" - log "github.com/sirupsen/logrus" "github.com/swuecho/chat_backend/auth" "github.com/swuecho/chat_backend/dto" "github.com/swuecho/chat_backend/sqlc_queries" + "log/slog" ) // --- Request types --- @@ -39,7 +39,7 @@ func (h *AuthUserHandler) ResetPasswordHandler(w http.ResponseWriter, r *http.Re tempPassword, err := auth.GenerateRandomPassword() if err != nil { - log.WithError(err).Error("Failed to generate temporary password") + slog.Error("Failed to generate temporary password", "error", err) dto.RespondWithAPIError(w, dto.ErrInternalUnexpected.WithMessage("Failed to generate temporary password").WithDebugInfo(err.Error())) return } diff --git a/api/handler/auth_user.go b/api/handler/auth_user.go index 82162189..70acb26b 100644 --- a/api/handler/auth_user.go +++ b/api/handler/auth_user.go @@ -9,12 +9,12 @@ import ( "time" "github.com/gorilla/mux" - log "github.com/sirupsen/logrus" "github.com/swuecho/chat_backend/auth" "github.com/swuecho/chat_backend/dto" "github.com/swuecho/chat_backend/middleware" "github.com/swuecho/chat_backend/sqlc_queries" "github.com/swuecho/chat_backend/svc" + "log/slog" ) // Token lifetime constants. @@ -38,7 +38,7 @@ type AuthUserHandler struct { // NewAuthUserHandler creates a new AuthUserHandler. func NewAuthUserHandler(sqlc_q *sqlc_queries.Queries, jwtSecret, audience string, defaultRateLimit int32) *AuthUserHandler { return &AuthUserHandler{ - service: svc.NewAuthUserService(sqlc_q), + service: svc.NewAuthUserService(sqlc_q, jwtSecret, defaultRateLimit), jwtSecret: jwtSecret, audience: audience, defaultRateLimit: defaultRateLimit, @@ -153,7 +153,7 @@ func (h *AuthUserHandler) UpdateSelf(w http.ResponseWriter, r *http.Request) { params.ID = userID user, err := h.service.UpdateAuthUser(r.Context(), params) if err != nil { - log.WithError(err).Error("Failed to update user") + slog.Error("Failed to update user", "error", err) dto.RespondWithAPIError(w, dto.ErrInternalUnexpected.WithMessage("Failed to update user").WithDebugInfo(err.Error())) return } @@ -184,14 +184,14 @@ type LoginParams struct { func (h *AuthUserHandler) SignUp(w http.ResponseWriter, r *http.Request) { var params LoginParams if err := json.NewDecoder(r.Body).Decode(¶ms); err != nil { - log.WithFields(log.Fields{"error": err, "ip": r.RemoteAddr, "action": "signup_decode_error"}).Warn("Failed to decode signup") + slog.Warn("Failed to decode signup", "error", err, "ip", r.RemoteAddr, "action", "signup_decode_error") dto.RespondWithAPIError(w, dto.ErrValidationInvalidInput("Invalid request: unable to decode JSON body").WithDebugInfo(err.Error())) return } hash, err := auth.GeneratePasswordHash(params.Password) if err != nil { - log.WithFields(log.Fields{"email": params.Email, "error": err}).Error("Failed to hash password") + slog.Error("Failed to hash password", "email", params.Email, "error", err) dto.RespondWithAPIError(w, dto.ErrInternalUnexpected.WithMessage("Failed to generate password hash").WithDebugInfo(err.Error())) return } @@ -200,7 +200,7 @@ func (h *AuthUserHandler) SignUp(w http.ResponseWriter, r *http.Request) { Password: hash, Email: params.Email, Username: params.Email, }) if err != nil { - log.WithFields(log.Fields{"email": params.Email, "error": err}).Error("Failed to create user") + slog.Error("Failed to create user", "email", params.Email, "error", err) dto.RespondWithAPIError(w, dto.WrapError(err, "Failed to create user")) return } @@ -220,7 +220,7 @@ func (h *AuthUserHandler) SignUp(w http.ResponseWriter, r *http.Request) { http.SetCookie(w, createSecureRefreshCookie(RefreshTokenName, refreshToken, int(refreshLifetime.Seconds()), r)) - log.WithFields(log.Fields{"user_id": user.ID, "email": user.Email, "action": "signup_success"}).Info("User signup successful") + slog.Info("User signup successful", "user_id", user.ID, "email", user.Email, "action", "signup_success") w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(dto.TokenResult{AccessToken: accessToken, ExpiresIn: int(time.Now().Add(AccessTokenLifetime).Unix())}) @@ -229,14 +229,14 @@ func (h *AuthUserHandler) SignUp(w http.ResponseWriter, r *http.Request) { func (h *AuthUserHandler) Login(w http.ResponseWriter, r *http.Request) { var params LoginParams if err := json.NewDecoder(r.Body).Decode(¶ms); err != nil { - log.WithFields(log.Fields{"error": err, "ip": r.RemoteAddr, "action": "login_decode_error"}).Warn("Failed to decode login") + slog.Warn("Failed to decode login", "error", err, "ip", r.RemoteAddr, "action", "login_decode_error") dto.RespondWithAPIError(w, dto.ErrValidationInvalidInput("Failed to decode request body").WithDebugInfo(err.Error())) return } user, err := h.service.Authenticate(r.Context(), params.Email, params.Password) if err != nil { - log.WithFields(log.Fields{"email": params.Email, "ip": r.RemoteAddr, "error": err, "action": "login_failed"}).Warn("User login failed") + slog.Warn("User login failed", "email", params.Email, "ip", r.RemoteAddr, "error", err, "action", "login_failed") dto.RespondWithAPIError(w, dto.ErrAuthInvalidEmailOrPassword.WithDebugInfo(err.Error())) return } @@ -256,7 +256,7 @@ func (h *AuthUserHandler) Login(w http.ResponseWriter, r *http.Request) { http.SetCookie(w, createSecureRefreshCookie(RefreshTokenName, refreshToken, int(refreshLifetime.Seconds()), r)) - log.WithFields(log.Fields{"user_id": user.ID, "email": user.Email, "action": "login_success"}).Info("User login successful") + slog.Info("User login successful", "user_id", user.ID, "email", user.Email, "action", "login_success") w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(dto.TokenResult{AccessToken: accessToken, ExpiresIn: int(time.Now().Add(AccessTokenLifetime).Unix())}) @@ -278,18 +278,18 @@ func (h *AuthUserHandler) ForeverToken(w http.ResponseWriter, r *http.Request) { } func (h *AuthUserHandler) RefreshToken(w http.ResponseWriter, r *http.Request) { - log.WithFields(log.Fields{"ip": r.RemoteAddr, "action": "refresh_attempt"}).Info("Token refresh attempt") + slog.Info("Token refresh attempt", "ip", r.RemoteAddr, "action", "refresh_attempt") refreshCookie, err := r.Cookie(RefreshTokenName) if err != nil { - log.WithFields(log.Fields{"ip": r.RemoteAddr, "error": err, "action": "refresh_missing_cookie"}).Warn("Missing refresh token cookie") + slog.Warn("Missing refresh token cookie", "ip", r.RemoteAddr, "error", err, "action", "refresh_missing_cookie") dto.RespondWithAPIError(w, dto.ErrAuthInvalidCredentials.WithMessage("Missing refresh token")) return } - result := middleware.ParseAndValidateJWT(refreshCookie.Value, auth.TokenTypeRefresh) + result := middleware.ParseAndValidateJWT(refreshCookie.Value, auth.TokenTypeRefresh, h.jwtSecret) if result.Error != nil { - log.WithFields(log.Fields{"ip": r.RemoteAddr, "error": result.Error.Detail, "action": "refresh_invalid_token"}).Warn("Invalid refresh token") + slog.Warn("Invalid refresh token", "ip", r.RemoteAddr, "error", result.Error.Detail, "action", "refresh_invalid_token") dto.RespondWithAPIError(w, *result.Error) return } @@ -306,7 +306,7 @@ func (h *AuthUserHandler) RefreshToken(w http.ResponseWriter, r *http.Request) { return } - log.WithFields(log.Fields{"user_id": userIDInt, "action": "refresh_success"}).Info("Token refresh successful") + slog.Info("Token refresh successful", "user_id", userIDInt, "action", "refresh_success") w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(dto.TokenResult{AccessToken: accessToken, ExpiresIn: int(time.Now().Add(AccessTokenLifetime).Unix())}) diff --git a/api/handler/chat_main.go b/api/handler/chat_main.go index e269ef07..5c6550da 100644 --- a/api/handler/chat_main.go +++ b/api/handler/chat_main.go @@ -1,14 +1,13 @@ package handler import ( - "context" "encoding/json" "net/http" "time" - log "github.com/sirupsen/logrus" "github.com/gorilla/mux" "golang.org/x/time/rate" + "log/slog" "github.com/swuecho/chat_backend/dto" "github.com/swuecho/chat_backend/provider" @@ -21,20 +20,22 @@ type ChatHandler struct { service *svc.ChatService sessionSvc *svc.ChatSessionService chatfileService *svc.ChatFileService - requestCtx context.Context rateLimiter *rate.Limiter + openAIKey string + openAIProxy string } const sessionTitleGenerationTimeout = 30 * time.Second // NewChatHandler creates a new ChatHandler. -func NewChatHandler(sqlc_q *sqlc_queries.Queries, rateLimiter *rate.Limiter) *ChatHandler { +func NewChatHandler(sqlc_q *sqlc_queries.Queries, rateLimiter *rate.Limiter, openAIKey, openAIProxy string) *ChatHandler { return &ChatHandler{ - service: svc.NewChatService(sqlc_q), + service: svc.NewChatService(sqlc_q, openAIKey, openAIProxy), sessionSvc: svc.NewChatSessionService(sqlc_q), chatfileService: svc.NewChatFileService(sqlc_q), - requestCtx: context.Background(), rateLimiter: rateLimiter, + openAIKey: openAIKey, + openAIProxy: openAIProxy, } } @@ -47,26 +48,20 @@ func (h *ChatHandler) Register(router *mux.Router) { // --- provider.Handler implementation --- -func (h *ChatHandler) RequestContext() context.Context { return h.requestCtx } -func (h *ChatHandler) Queries() *sqlc_queries.Queries { return h.service.Q() } +func (h *ChatHandler) Queries() *sqlc_queries.Queries { return h.service.Q() } func (h *ChatHandler) Config() provider.Config { return provider.Config{ - OpenAIKey: svc.Cfg.OpenAIKey, - OpenAIProxy: svc.Cfg.OpenAIProxy, + OpenAIKey: h.openAIKey, + OpenAIProxy: h.openAIProxy, RateLimiter: h.rateLimiter, } } -// GetRequestContext is a legacy alias for RequestContext. -func (h *ChatHandler) GetRequestContext() context.Context { - return h.requestCtx -} - // GetChatInstructions returns artifact instruction text. func (h *ChatHandler) GetChatInstructions(w http.ResponseWriter, r *http.Request) { artifactInstruction, err := svc.LoadArtifactInstruction() if err != nil { - log.Printf("Warning: Failed to load artifact instruction: %v", err) + slog.Warn("Failed to load artifact instruction", "error", err) artifactInstruction = "" } json.NewEncoder(w).Encode(dto.ChatInstructionResponse{ diff --git a/api/handler/chat_model_select.go b/api/handler/chat_model_select.go index b756dbad..8941af70 100644 --- a/api/handler/chat_model_select.go +++ b/api/handler/chat_model_select.go @@ -1,13 +1,13 @@ package handler import ( + "context" "database/sql" "errors" "fmt" - "net/http" - log "github.com/sirupsen/logrus" mapset "github.com/deckarep/golang-set/v2" + "log/slog" "github.com/swuecho/chat_backend/dto" "github.com/swuecho/chat_backend/models" @@ -16,12 +16,12 @@ import ( ) // chooseChatModel returns the appropriate ChatModel implementation based on session config. -func (h *ChatHandler) chooseChatModel(session sqlc_queries.ChatSession, msgs []models.Message) provider.ChatModel { +func (h *ChatHandler) chooseChatModel(ctx context.Context, session sqlc_queries.ChatSession, msgs []models.Message) provider.ChatModel { if isTest(msgs) { return provider.NewTestChatModel(h) } - chatModel, err := provider.GetChatModel(h.RequestContext(), h.Queries(), session.Model) + chatModel, err := provider.GetChatModel(ctx, h.Queries(), session.Model) if err != nil { return provider.NewOpenAIChatModel(h) // fallback } @@ -59,18 +59,17 @@ func isTest(msgs []models.Message) bool { } // CheckModelAccess verifies the user hasn't exceeded per-model rate limits. -func (h *ChatHandler) CheckModelAccess(w http.ResponseWriter, chatSessionUuid, model string, userID int32) bool { - ctx := h.RequestContext() - +// Returns nil if access is allowed, or an error (dto.APIError) if denied. +func (h *ChatHandler) CheckModelAccess(ctx context.Context, chatSessionUuid, model string, userID int32) error { chatModel, err := h.sessionSvc.ChatModelByName(ctx, model) if err != nil { - log.WithError(err).WithField("model", model).Error("Chat model not found") - dto.RespondWithAPIError(w, dto.ErrResourceNotFound("chat model: "+model)) - return true + slog.Error("Chat model not found", "error", err, "model", model) + apiErr := dto.ErrResourceNotFound("chat model: " + model) + return apiErr } if !chatModel.EnablePerModeRatelimit { - return false + return nil } rate, err := h.sessionSvc.RateLimitByUserAndSessionUUID(ctx, sqlc_queries.RateLimiteByUserAndSessionUUIDParams{ @@ -78,27 +77,24 @@ func (h *ChatHandler) CheckModelAccess(w http.ResponseWriter, chatSessionUuid, m }) if err != nil { if errors.Is(err, sql.ErrNoRows) { - return false + return nil } - dto.RespondWithAPIError(w, dto.WrapError(dto.MapDatabaseError(err), "Failed to get rate limit")) - return true + return dto.WrapError(dto.MapDatabaseError(err), "Failed to get rate limit") } usage10Min, err := h.sessionSvc.GetChatMessagesCountByUserAndModel(ctx, sqlc_queries.GetChatMessagesCountByUserAndModelParams{ UserID: userID, Model: rate.ChatModelName, }) if err != nil { - dto.RespondWithAPIError(w, dto.ErrInternalUnexpected.WithDetail("Failed to get usage data").WithDebugInfo(err.Error())) - return true + return dto.ErrInternalUnexpected.WithDetail("Failed to get usage data").WithDebugInfo(err.Error()) } if int32(usage10Min) > rate.RateLimit { apiErr := dto.ErrTooManyRequests apiErr.Message = fmt.Sprintf("Rate limit exceeded for %s", rate.ChatModelName) apiErr.Detail = fmt.Sprintf("Usage: %d, Limit: %d", usage10Min, rate.RateLimit) - dto.RespondWithAPIError(w, apiErr) - return true + return apiErr } - return false + return nil } diff --git a/api/handler/chat_session_helpers.go b/api/handler/chat_session_helpers.go index f23dd072..58a65685 100644 --- a/api/handler/chat_session_helpers.go +++ b/api/handler/chat_session_helpers.go @@ -9,9 +9,11 @@ import ( "net/http" "strings" - log "github.com/sirupsen/logrus" + openai "github.com/sashabaranov/go-openai" + "log/slog" "github.com/swuecho/chat_backend/dto" + "github.com/swuecho/chat_backend/models" "github.com/swuecho/chat_backend/provider" "github.com/swuecho/chat_backend/sqlc_queries" ) @@ -20,7 +22,7 @@ import ( func (h *ChatHandler) validateChatSession(ctx context.Context, w http.ResponseWriter, chatSessionUuid string) (*sqlc_queries.ChatSession, *sqlc_queries.ChatModel, string, bool) { chatSession, err := h.sessionSvc.GetChatSessionByUUID(ctx, chatSessionUuid) if err != nil { - log.Printf("Invalid session UUID: %s, error: %v", chatSessionUuid, err) + slog.Info("Invalid session UUID", "uuid", chatSessionUuid, "error", err) dto.RespondWithAPIError(w, dto.ErrResourceNotFound("chat session").WithMessage(chatSessionUuid)) return nil, nil, "", false } @@ -49,7 +51,7 @@ func (h *ChatHandler) handlePromptCreation(ctx context.Context, w http.ResponseW if errors.Is(err, sql.ErrNoRows) { existingPrompt = false } else { - log.Printf("Error checking prompt for session %s: %v", chatSession.Uuid, err) + slog.Error("error checking prompt", "session", chatSession.Uuid, "error", err) dto.RespondWithAPIError(w, dto.CreateAPIError(dto.ErrInternalUnexpected, "Failed to get prompt", err.Error())) return false } @@ -79,7 +81,7 @@ func (h *ChatHandler) handlePromptCreation(ctx context.Context, w http.ResponseW Uuid: chatSession.Uuid, UserID: userID, Topic: title, } if _, err := h.sessionSvc.UpdateChatSessionTopicByUUID(ctx, params); err != nil { - log.Printf("Warning: Failed to update session title: %v", err) + slog.Warn("Failed to update session title", "error", err) } } } @@ -91,17 +93,16 @@ func (h *ChatHandler) handlePromptCreation(ctx context.Context, w http.ResponseW func (h *ChatHandler) generateAndSaveAnswer(ctx context.Context, w http.ResponseWriter, chatSession *sqlc_queries.ChatSession, chatUuid string, userID int32, baseURL string, streamOutput bool) bool { msgs, err := h.service.GetAskMessages(*chatSession, chatUuid, false) if err != nil { - log.Printf("Error collecting messages for session %s: %v", chatSession.Uuid, err) + slog.Error("error collecting messages", "session", chatSession.Uuid, "error", err) dto.RespondWithAPIError(w, dto.CreateAPIError(dto.ErrInternalUnexpected, "Failed to collect messages", err.Error())) return false } - log.Printf("Collected messages - SessionUUID: %s, Count: %d, Model: %s", chatSession.Uuid, len(msgs), chatSession.Model) + slog.Info("Collected messages", "sessionUUID", chatSession.Uuid, "count", len(msgs), "model", chatSession.Model) - h.requestCtx = ctx - model := h.chooseChatModel(*chatSession, msgs) - LLMAnswer, err := model.Stream(w, *chatSession, msgs, chatUuid, false, streamOutput) + model := h.chooseChatModel(ctx, *chatSession, msgs) + LLMAnswer, err := streamFromModel(model, ctx, w, *chatSession, msgs, chatUuid, false, streamOutput) if err != nil { - log.Printf("Error generating answer: %v", err) + slog.Error("error generating answer", "error", err) dto.RespondWithAPIError(w, dto.WrapError(err, "Failed to generate answer")) return false } @@ -128,6 +129,62 @@ func (h *ChatHandler) generateAndSaveAnswer(ctx context.Context, w http.Response return true } +// streamFromModel calls model.Stream() and consumes the channel, writing SSE or JSON to w. +// Returns the final answer or an error. +func streamFromModel(model provider.ChatModel, ctx context.Context, w http.ResponseWriter, session sqlc_queries.ChatSession, msgs []models.Message, chatUuid string, regenerate bool, streamOutput bool) (*models.LLMAnswer, error) { + ch, err := model.Stream(ctx, session, msgs, chatUuid, regenerate, streamOutput) + if err != nil { + return nil, err + } + + var lastAnswer *models.LLMAnswer + + if streamOutput { + flusher, err := setupSSEStream(w) + if err != nil { + return nil, err + } + for chunk := range ch { + if chunk.Err != nil { + return nil, chunk.Err + } + if chunk.Done { + lastAnswer = chunk.FinalAnswer + break + } + if chunk.Content != "" { + provider.FlushResponse(w, flusher, provider.StreamingResponse{ + AnswerID: chunk.ID, + Content: chunk.Content, + IsFinal: false, + }) + } + } + } else { + for chunk := range ch { + if chunk.Err != nil { + return nil, chunk.Err + } + if chunk.Done { + lastAnswer = chunk.FinalAnswer + break + } + } + // Write non-streaming JSON response + if lastAnswer != nil { + json.NewEncoder(w).Encode(ChatCompletionResponse{ + ID: lastAnswer.AnswerId, + Object: "chat.completion", + Choices: []Choice{{ + Message: openai.ChatCompletionMessage{Content: lastAnswer.Answer}, + }}, + }) + } + } + + return lastAnswer, nil +} + // generateSessionTitle asynchronously updates the session topic using an LLM. func (h *ChatHandler) generateSessionTitle(chatSession *sqlc_queries.ChatSession, userID int32) { ctx, cancel := context.WithTimeout(context.Background(), sessionTitleGenerationTimeout) @@ -137,7 +194,7 @@ func (h *ChatHandler) generateSessionTitle(chatSession *sqlc_queries.ChatSession Uuid: chatSession.Uuid, Offset: 0, Limit: 100, }) if err != nil { - log.Printf("Warning: Failed to get messages for title generation: %v", err) + slog.Warn("Failed to get messages for title generation", "error", err) return } @@ -163,11 +220,11 @@ func (h *ChatHandler) generateSessionTitle(chatSession *sqlc_queries.ChatSession if _, err := h.sessionSvc.UpdateChatSessionTopicByUUID(ctx, sqlc_queries.UpdateChatSessionTopicByUUIDParams{ Uuid: chatSession.Uuid, UserID: userID, Topic: genTitle, }); err != nil { - log.Printf("Warning: Failed to update session title: %v", err) + slog.Warn("Failed to update session title", "error", err) return } - log.Printf("Generated LLM title for session %s: %s", chatSession.Uuid, genTitle) + slog.Info("Generated LLM title", "sessionUUID", chatSession.Uuid, "title", genTitle) } // sendSuggestedQuestionsStream sends suggested questions as an SSE event. diff --git a/api/handler/chat_stream.go b/api/handler/chat_stream.go index ab1c9b1f..6ae8e278 100644 --- a/api/handler/chat_stream.go +++ b/api/handler/chat_stream.go @@ -6,8 +6,8 @@ import ( "fmt" "net/http" - log "github.com/sirupsen/logrus" openai "github.com/sashabaranov/go-openai" + "log/slog" "github.com/swuecho/chat_backend/dto" "github.com/swuecho/chat_backend/models" @@ -31,10 +31,10 @@ type BotRequest struct { } type ChatCompletionResponse struct { - ID string `json:"id"` - Object string `json:"object"` - Created int `json:"created"` - Model string `json:"model"` + ID string `json:"id"` + Object string `json:"object"` + Created int `json:"created"` + Model string `json:"model"` Usage struct { PromptTokens int `json:"prompt_tokens"` CompletionTokens int `json:"completion_tokens"` @@ -62,7 +62,7 @@ func (h *ChatHandler) ChatBotCompletionHandler(w http.ResponseWriter, r *http.Re ctx := r.Context() userID, err := getUserID(ctx) if err != nil { - log.Printf("Error getting user ID: %v", err) + slog.Error("error getting user ID", "error", err) dto.RespondWithAPIError(w, dto.ErrAuthInvalidCredentials.WithDebugInfo(err.Error())) return } @@ -94,7 +94,7 @@ func (h *ChatHandler) ChatBotCompletionHandler(w http.ResponseWriter, r *http.Re func (h *ChatHandler) ChatCompletionHandler(w http.ResponseWriter, r *http.Request) { var req ChatRequest if err := json.NewDecoder(r.Body).Decode(&req); err != nil { - log.Printf("Error decoding request: %v", err) + slog.Error("error decoding request", "error", err) dto.RespondWithAPIError(w, dto.ErrValidationInvalidInput("Invalid request format").WithDebugInfo(err.Error())) return } @@ -102,7 +102,7 @@ func (h *ChatHandler) ChatCompletionHandler(w http.ResponseWriter, r *http.Reque ctx := r.Context() userID, err := getUserID(ctx) if err != nil { - log.Printf("Error getting user ID: %v", err) + slog.Error("error getting user ID", "error", err) dto.RespondWithAPIError(w, dto.ErrAuthInvalidCredentials.WithDebugInfo(err.Error())) return } @@ -120,7 +120,7 @@ func genAnswer(h *ChatHandler, w http.ResponseWriter, ctx context.Context, sessi if !ok { return } - log.Printf("Processing chat session - SessionUUID: %s, UserID: %d, Model: %s", chatSession.Uuid, userID, chatSession.Model) + slog.Info("Processing chat session", "sessionUUID", chatSession.Uuid, "userID", userID, "model", chatSession.Model) if !h.handlePromptCreation(ctx, w, chatSession, chatUuid, question, userID, baseURL) { return @@ -131,7 +131,7 @@ func genAnswer(h *ChatHandler, w http.ResponseWriter, ctx context.Context, sessi // genBotAnswer generates a bot answer from a snapshot conversation. func genBotAnswer(h *ChatHandler, w http.ResponseWriter, session sqlc_queries.ChatSession, messages []dto.SimpleChatMessage, snapshotUuid, question string, userID int32, streamOutput bool) { - ctx := h.GetRequestContext() + ctx := context.Background() if _, err := h.sessionSvc.ChatModelByName(ctx, session.Model); err != nil { dto.RespondWithAPIError(w, dto.ErrResourceNotFound("Chat model: "+session.Model).WithDebugInfo(err.Error())) return @@ -140,8 +140,8 @@ func genBotAnswer(h *ChatHandler, w http.ResponseWriter, session sqlc_queries.Ch msgs := simpleChatMessagesToMessages(messages) msgs = append(msgs, models.Message{Role: "user", Content: question}) - model := h.chooseChatModel(session, msgs) - LLMAnswer, err := model.Stream(w, session, msgs, "", false, streamOutput) + model := h.chooseChatModel(ctx, session, msgs) + LLMAnswer, err := streamFromModel(model, ctx, w, session, msgs, "", false, streamOutput) if err != nil { dto.RespondWithAPIError(w, dto.WrapError(err, "Failed to generate answer")) return @@ -155,7 +155,7 @@ func genBotAnswer(h *ChatHandler, w http.ResponseWriter, session sqlc_queries.Ch Model: session.Model, TokensUsed: int32(len(LLMAnswer.Answer)) / 4, }); err != nil { - log.Printf("Failed to save bot answer history: %v", err) + slog.Info("Failed to save bot answer history", "error", err) } if !isTest(msgs) { @@ -176,11 +176,11 @@ func regenerateAnswer(h *ChatHandler, w http.ResponseWriter, ctx context.Context return } - h.requestCtx = ctx - model := h.chooseChatModel(*chatSession, msgs) - LLMAnswer, err := model.Stream(w, *chatSession, msgs, chatUuid, true, stream) + model := h.chooseChatModel(ctx, *chatSession, msgs) + LLMAnswer, err := streamFromModel(model, ctx, w, *chatSession, msgs, chatUuid, true, stream) if err != nil { - log.Printf("Error regenerating answer: %v", err) + slog.Error("error regenerating answer", "error", err) + dto.RespondWithAPIError(w, dto.WrapError(err, "Failed to regenerate answer")) return } diff --git a/api/handler/file_upload.go b/api/handler/file_upload.go index 93aa6df0..d06653a5 100644 --- a/api/handler/file_upload.go +++ b/api/handler/file_upload.go @@ -7,7 +7,7 @@ import ( "errors" "fmt" "io" - "log" + "log/slog" "net/http" "strconv" @@ -116,7 +116,7 @@ func (h *ChatFileHandler) DownloadFile(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%s", file.Name)) w.Header().Set("Content-Length", strconv.Itoa(len(file.Data))) if _, err := w.Write(file.Data); err != nil { - log.Printf("Failed to write file data: %v", err) + slog.Info("Failed to write file data", "error", err) } } diff --git a/api/handler/message.go b/api/handler/message.go index 9d8ec74d..d68c3b06 100644 --- a/api/handler/message.go +++ b/api/handler/message.go @@ -17,14 +17,18 @@ import ( ) type ChatMessageHandler struct { - service *svc.ChatMessageService - sessionSvc *svc.ChatSessionService + service *svc.ChatMessageService + sessionSvc *svc.ChatSessionService + openAIKey string + openAIProxy string } -func NewChatMessageHandler(sqlc_q *sqlc_queries.Queries) *ChatMessageHandler { +func NewChatMessageHandler(sqlc_q *sqlc_queries.Queries, openAIKey, openAIProxy string) *ChatMessageHandler { return &ChatMessageHandler{ - service: svc.NewChatMessageService(sqlc_q), - sessionSvc: svc.NewChatSessionService(sqlc_q), + service: svc.NewChatMessageService(sqlc_q), + sessionSvc: svc.NewChatSessionService(sqlc_q), + openAIKey: openAIKey, + openAIProxy: openAIProxy, } } @@ -273,7 +277,7 @@ func (h *ChatMessageHandler) GenerateMoreSuggestions(w http.ResponseWriter, r *h }) } - chatService := svc.NewChatService(h.service.Q()) + chatService := svc.NewChatService(h.service.Q(), h.openAIKey, h.openAIProxy) newSuggestions := chatService.GenerateSuggestedQuestions(message.Content, msgs) if len(newSuggestions) == 0 { diff --git a/api/handler/model_privilege.go b/api/handler/model_privilege.go index 7de27f7a..3dffd2da 100644 --- a/api/handler/model_privilege.go +++ b/api/handler/model_privilege.go @@ -5,7 +5,7 @@ import ( "encoding/json" "errors" "fmt" - "log" + "log/slog" "net/http" "strconv" @@ -80,7 +80,7 @@ func (h *UserChatModelPrivilegeHandler) CreateUserChatModelPrivilege(w http.Resp return } - log.Printf("Creating chat model privilege for user %s with model %s", input.UserEmail, input.ChatModelName) + slog.Info("Creating chat model privilege for user %s with model %s", input.UserEmail, input.ChatModelName) user, err := h.db.GetAuthUserByEmail(r.Context(), input.UserEmail) if err != nil { diff --git a/api/handler/user_active_session.go b/api/handler/user_active_session.go index f3202a12..9ea0a38c 100644 --- a/api/handler/user_active_session.go +++ b/api/handler/user_active_session.go @@ -5,7 +5,7 @@ import ( "encoding/json" "errors" "fmt" - "log" + "log/slog" "net/http" "regexp" @@ -55,7 +55,7 @@ func (h *UserActiveChatSessionHandler) GetUserActiveChatSessionHandler(w http.Re w.Header().Set("Content-Type", "application/json") if err := json.NewEncoder(w).Encode(session); err != nil { - log.Printf("Failed to encode response: %v", err) + slog.Info("Failed to encode response", "error", err) } } @@ -81,7 +81,7 @@ func (h *UserActiveChatSessionHandler) CreateOrUpdateUserActiveChatSessionHandle return } - log.Printf("Creating/updating active chat session for user %d", userID) + slog.Info("Creating/updating active chat session", "userID", userID) session, err := h.service.UpsertActiveSession(r.Context(), userID, nil, reqBody.ChatSessionUuid) if err != nil { @@ -91,7 +91,7 @@ func (h *UserActiveChatSessionHandler) CreateOrUpdateUserActiveChatSessionHandle w.Header().Set("Content-Type", "application/json") if err := json.NewEncoder(w).Encode(session); err != nil { - log.Printf("Failed to encode response: %v", err) + slog.Info("Failed to encode response", "error", err) } } diff --git a/api/handler/util.go b/api/handler/util.go index b6634136..810d6f8f 100644 --- a/api/handler/util.go +++ b/api/handler/util.go @@ -1,165 +1,29 @@ // Package handler provides HTTP request handlers for the chat API. -// -// Handlers are organized by domain: -// - Auth: authentication, user management -// - Chat: streaming, completion, session/message/prompt CRUD -// - Workspace: workspace and session organization -// - Model: LLM model management and privileges -// - Snapshot: conversation snapshots and bots -// - File: file upload and download -// - Admin: user stats, rate limits package handler import ( "context" - "encoding/json" - "errors" - "fmt" "net/http" - "os" - "strconv" - "github.com/google/uuid" - "github.com/pkoukk/tiktoken-go" - "github.com/swuecho/chat_backend/dto" - "github.com/swuecho/chat_backend/middleware" + "github.com/swuecho/chat_backend/pkg/util" ) -// NewUUID generates a new UUID v7 string. -func NewUUID() string { - uuidv7, err := uuid.NewV7() - if err != nil { - return uuid.NewString() - } - return uuidv7.String() -} +// Re-exported from pkg/util for convenient use within the handler package. -// getTokenCount returns the estimated token count for a text string. -func getTokenCount(content string) (int, error) { - tke, err := tiktoken.GetEncoding("cl100k_base") - if err != nil { - return 0, err - } - return len(tke.Encode(content, nil, nil)), nil -} - -// firstNWords returns the first n words of a string. -func firstNWords(s string, n int) string { - if s == "" { - return "" - } - words := splitWords(s) - if len(words) <= n { - return s - } - result := "" - for i, w := range words[:n] { - if i > 0 { - result += " " - } - result += w - } - return result -} - -// splitWords splits a string by whitespace. -func splitWords(s string) []string { - var words []string - word := "" - for _, r := range s { - if r == ' ' || r == '\t' || r == '\n' { - if word != "" { - words = append(words, word) - word = "" - } - } else { - word += string(r) - } - } - if word != "" { - words = append(words, word) - } - return words -} - -// setSSEHeader sets headers for Server-Sent Events. -func setSSEHeader(w http.ResponseWriter) { - w.Header().Set("Content-Type", "text/event-stream") - w.Header().Set("Cache-Control", "no-cache, no-transform") - w.Header().Set("Connection", "keep-alive") - w.Header().Set("Access-Control-Allow-Origin", "*") - w.Header().Set("X-Accel-Buffering", "no") - w.Header().Del("Content-Length") - w.Header().Del("Content-Encoding") -} - -// setupSSEStream prepares an SSE stream and returns a Flusher. -func setupSSEStream(w http.ResponseWriter) (http.Flusher, error) { - setSSEHeader(w) - flusher, ok := w.(http.Flusher) - if !ok { - return nil, errors.New(dto.ErrorStreamUnsupported) - } - return flusher, nil -} - -// getPerWordStreamLimit returns the per-word streaming limit from env. -func getPerWordStreamLimit() int { - val := os.Getenv("PER_WORD_STREAM_LIMIT") - if val == "" { - return 200 - } - n, err := strconv.Atoi(val) - if err != nil { - return 200 - } - return n -} - -// getPaginationParams extracts limit and offset from query parameters. -func getPaginationParams(r *http.Request) (limit int32, offset int32) { - limit = 100 - if v := r.URL.Query().Get("limit"); v != "" { - if l, err := strconv.ParseInt(v, 10, 32); err == nil { - limit = int32(l) - } - } - if v := r.URL.Query().Get("offset"); v != "" { - if o, err := strconv.ParseInt(v, 10, 32); err == nil { - offset = int32(o) - } - } - return -} - -// getLimitParam extracts an optional limit parameter with a default. -func getLimitParam(r *http.Request, defaultLimit int32) int32 { - v := r.URL.Query().Get("limit") - if v == "" { - return defaultLimit - } - n, err := strconv.ParseInt(v, 10, 32) - if err != nil { - return defaultLimit - } - return int32(n) -} - -// DecodeJSON decodes a JSON request body into the target. -func DecodeJSON(r *http.Request, target any) error { - return json.NewDecoder(r.Body).Decode(target) -} +var ( + NewUUID = util.NewUUID + SetupSSE = util.SetupSSE + PerWordStreamLimit = util.PerWordStreamLimit + PaginationParams = util.PaginationParams + LimitParam = util.LimitParam + DecodeJSON = util.DecodeJSON +) -// getUserID extracts the authenticated user ID from the request context. -func getUserID(ctx context.Context) (int32, error) { - userIdValue := ctx.Value(middleware.UserContextKey) - if userIdValue == nil { - return 0, fmt.Errorf("no user Id in context") - } - userIDStr, _ := userIdValue.(string) - userIDInt, err := strconv.ParseInt(userIDStr, 10, 32) - if err != nil { - return 0, fmt.Errorf("invalid user ID: %s", userIDStr) - } - return int32(userIDInt), nil -} +func getTokenCount(content string) (int, error) { return util.TokenCount(content) } +func firstNWords(s string, n int) string { return util.FirstNWords(s, n) } +func getUserID(ctx context.Context) (int32, error) { return util.UserID(ctx) } +func setSSEHeader(w http.ResponseWriter) { _, _ = util.SetupSSE(w) } +func setupSSEStream(w http.ResponseWriter) (http.Flusher, error) { return util.SetupSSE(w) } +func getPerWordStreamLimit() int { return util.PerWordStreamLimit() } +func getPaginationParams(r *http.Request) (int32, int32) { return util.PaginationParams(r) } +func getLimitParam(r *http.Request, d int32) int32 { return util.LimitParam(r, d) } diff --git a/api/handler/workspace_crud.go b/api/handler/workspace_crud.go index 98fed36d..0b384665 100644 --- a/api/handler/workspace_crud.go +++ b/api/handler/workspace_crud.go @@ -4,7 +4,7 @@ import ( "context" "database/sql" "encoding/json" - "log" + "log/slog" "net/http" "github.com/google/uuid" @@ -61,7 +61,7 @@ func (h *ChatWorkspaceHandler) createWorkspace(w http.ResponseWriter, r *http.Re func (h *ChatWorkspaceHandler) getWorkspaceByUUID(w http.ResponseWriter, r *http.Request) { workspaceUUID := mux.Vars(r)["uuid"] - log.Printf("getWorkspaceByUUID called with UUID=%s", workspaceUUID) + slog.Info("getWorkspaceByUUID called", "uuid", workspaceUUID) ctx := r.Context() userID, err := getUserID(ctx) diff --git a/api/handler/workspace_settings.go b/api/handler/workspace_settings.go index c0bfaaac..8c20303e 100644 --- a/api/handler/workspace_settings.go +++ b/api/handler/workspace_settings.go @@ -2,7 +2,7 @@ package handler import ( "encoding/json" - "log" + "log/slog" "net/http" "github.com/gorilla/mux" @@ -100,7 +100,7 @@ func (h *ChatWorkspaceHandler) autoMigrateLegacySessions(w http.ResponseWriter, if result.HasLegacySessions { if err := h.wsService.MigrateLegacyActiveSessions(ctx, userID, result.DefaultWorkspace.ID); err != nil { - log.Printf("Warning: failed to migrate legacy active sessions: %v", err) + slog.Warn("failed to migrate legacy active sessions", "error", err) } response["defaultWorkspace"] = workspaceToResponse(result.DefaultWorkspace) } diff --git a/api/llm/gemini/gemini.go b/api/llm/gemini/gemini.go index 9e979b88..bf0aaaa5 100644 --- a/api/llm/gemini/gemini.go +++ b/api/llm/gemini/gemini.go @@ -5,7 +5,7 @@ import ( "encoding/json" "fmt" "io" - "log" + "log/slog" "net/http" "os" "strings" @@ -210,12 +210,12 @@ func GenGemminPayload(chat_compeletion_messages []models.Message, chatFiles []sq } payloadBytes, err := json.Marshal(payload) - log.Printf("\n%s\n", string(payloadBytes)) if err != nil { fmt.Println("Error marshalling payload:", err) // handle err return nil, err } + slog.Info("gemini payload", "payload", string(payloadBytes)) return payloadBytes, nil } diff --git a/api/llm/openai/openai.go b/api/llm/openai/openai.go index 6c4ceb72..bdc71287 100644 --- a/api/llm/openai/openai.go +++ b/api/llm/openai/openai.go @@ -50,7 +50,7 @@ type ChatCompletionStreamResponse struct { SystemFingerprint string `json:"system_fingerprint"` PromptAnnotations []PromptAnnotation `json:"prompt_annotations,omitempty"` PromptFilterResults []PromptFilterResult `json:"prompt_filter_results,omitempty"` - // An optional field that will only be present when you set stream_options: {"include_usage": true} in your request. + // An optional field that will only be present when you set stream_options: {"include_usage", true} in your request. // When present, it contains a null value except for the last chunk which contains the token usage statistics // for the entire request. Usage *Usage `json:"usage,omitempty"` diff --git a/api/main.go b/api/main.go index 9d86a60b..accc14ba 100644 --- a/api/main.go +++ b/api/main.go @@ -20,7 +20,6 @@ import ( "github.com/swuecho/chat_backend/dto" "github.com/swuecho/chat_backend/handler" "github.com/swuecho/chat_backend/middleware" - "github.com/swuecho/chat_backend/provider" "github.com/swuecho/chat_backend/sqlc_queries" "github.com/swuecho/chat_backend/static" "github.com/swuecho/chat_backend/svc" @@ -49,15 +48,7 @@ func main() { func run() error { // --- Configuration --- - appCfg = config.Load() - cfg := appCfg - - // Initialize package-level configs - provider.OpenAIToken = cfg.OPENAI.API_KEY - svc.Cfg.OpenAIKey = cfg.OPENAI.API_KEY - svc.Cfg.OpenAIProxy = cfg.OPENAI.PROXY_URL - svc.Cfg.DefaultLimit = int32(cfg.OPENAI.RATELIMIT) - appCfg = cfg + cfg := config.Load() // --- Database --- pgdb, err := openDB(cfg) @@ -88,8 +79,6 @@ func run() error { return fmt.Errorf("jwt secret: %w", err) } srv.jwtSecret = jwtSecretAndAud - middleware.SetJWTSecret(jwtSecretAndAud.Secret) - svc.Cfg.JWTSecret = jwtSecretAndAud.Secret // --- Router --- router, rawRouter := srv.buildRouter() @@ -102,8 +91,8 @@ func run() error { // --- HTTP Server --- httpServer := &http.Server{ - Addr: ":8080", - Handler: router, + Addr: ":8080", + Handler: router, ReadTimeout: 60 * time.Second, ReadHeaderTimeout: 30 * time.Second, WriteTimeout: 10 * time.Minute, // Long write timeout for streaming @@ -179,8 +168,8 @@ func (s *server) buildRouter() (http.Handler, *mux.Router) { userRouter := apiRouter.NewRoute().Subrouter() // Auth middleware - adminRouter.Use(middleware.AdminAuthMiddleware) - userRouter.Use(middleware.UserAuthMiddleware) + adminRouter.Use(middleware.AdminAuthMiddleware(s.jwtSecret.Secret)) + userRouter.Use(middleware.UserAuthMiddleware(s.jwtSecret.Secret)) // Rate limiting rateLimitMW := middleware.RateLimitByUserID(s.q, s.cfg.OPENAI.RATELIMIT) @@ -213,6 +202,10 @@ func (s *server) buildRouter() (http.Handler, *mux.Router) { func (s *server) registerRoutes(apiRouter, adminRouter, userRouter *mux.Router) { q := s.q rateLimit := int32(s.cfg.OPENAI.RATELIMIT) + openAIKey := s.cfg.OPENAI.API_KEY + openAIProxy := s.cfg.OPENAI.PROXY_URL + jwtSecret := s.jwtSecret.Secret + jwtAudience := s.jwtSecret.Audience // Public apiRouter.HandleFunc("/tts", handler.HandleTTSRequest) @@ -222,12 +215,12 @@ func (s *server) registerRoutes(apiRouter, adminRouter, userRouter *mux.Router) handler.NewChatModelHandler(q).Register(userRouter) // Auth - authHandler := handler.NewAuthUserHandler(q, s.jwtSecret.Secret, s.jwtSecret.Audience, rateLimit) + authHandler := handler.NewAuthUserHandler(q, jwtSecret, jwtAudience, rateLimit) authHandler.Register(userRouter) authHandler.RegisterPublicRoutes(apiRouter) // Admin - handler.NewAdminHandler(svc.NewAuthUserService(q), rateLimit).RegisterRoutes(adminRouter) + handler.NewAdminHandler(svc.NewAuthUserService(q, jwtSecret, rateLimit), rateLimit).RegisterRoutes(adminRouter) // Prompts handler.NewChatPromptHandler(q).Register(userRouter) @@ -242,13 +235,13 @@ func (s *server) registerRoutes(apiRouter, adminRouter, userRouter *mux.Router) handler.NewChatWorkspaceHandler(q).Register(userRouter) // Messages - handler.NewChatMessageHandler(q).Register(userRouter) + handler.NewChatMessageHandler(q, openAIKey, openAIProxy).Register(userRouter) // Snapshots handler.NewChatSnapshotHandler(q).Register(userRouter) // Chat stream - handler.NewChatHandler(q, s.rateLimiter).Register(userRouter) + handler.NewChatHandler(q, s.rateLimiter, openAIKey, openAIProxy).Register(userRouter) // Model privileges handler.NewUserChatModelPrivilegeHandler(q).Register(userRouter) @@ -286,8 +279,8 @@ func (s *server) healthCheck(w http.ResponseWriter, r *http.Request) { } dto.RespondWithJSON(w, status, map[string]interface{}{ - "status": statusToText(healthy), - "checks": checks, + "status": statusToText(healthy), + "checks": checks, }) } @@ -353,7 +346,3 @@ func (s *server) idleMonitor() { } } } - -// Package-level var for the application config, used where dependency injection -// is not yet feasible (e.g., during initialization order). -var appCfg config.AppConfig diff --git a/api/middleware/auth.go b/api/middleware/auth.go index bd9f50e8..00b533fe 100644 --- a/api/middleware/auth.go +++ b/api/middleware/auth.go @@ -16,8 +16,8 @@ import ( "runtime/debug" "strings" - "github.com/google/uuid" jwt "github.com/golang-jwt/jwt/v5" + "github.com/google/uuid" "github.com/swuecho/chat_backend/auth" "github.com/swuecho/chat_backend/dto" ) @@ -31,14 +31,6 @@ const ( GuidContextKey contextKey = "guid" ) -// JWTSecret holds the JWT signing secret used across middleware. -var JWTSecret string - -// SetJWTSecret sets the JWT secret for middleware. -func SetJWTSecret(secret string) { - JWTSecret = secret -} - // ExtractBearerToken extracts the bearer token from an Authorization header. func ExtractBearerToken(r *http.Request) string { bearerToken := r.Header.Get("Authorization") @@ -57,7 +49,7 @@ func CreateUserContext(r *http.Request, userID, role string) *http.Request { } // ParseAndValidateJWT parses and validates a JWT token. -func ParseAndValidateJWT(bearerToken string, expectedTokenType string) *AuthTokenResult { +func ParseAndValidateJWT(bearerToken string, expectedTokenType string, jwtSecret string) *AuthTokenResult { result := &AuthTokenResult{} if bearerToken == "" { @@ -67,7 +59,7 @@ func ParseAndValidateJWT(bearerToken string, expectedTokenType string) *AuthToke return result } - jwtSigningKey := []byte(JWTSecret) + jwtSigningKey := []byte(jwtSecret) token, err := jwt.Parse(bearerToken, func(token *jwt.Token) (any, error) { if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok { return nil, fmt.Errorf("invalid JWT signing method") @@ -172,40 +164,44 @@ func parseInt32(s string) (int32, error) { } // AdminAuthMiddleware provides authentication + admin authorization. -func AdminAuthMiddleware(next http.Handler) http.Handler { - return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - bearerToken := ExtractBearerToken(r) - result := ParseAndValidateJWT(bearerToken, auth.TokenTypeAccess) - - if result.Error != nil { - dto.RespondWithAPIError(w, *result.Error) - return - } +func AdminAuthMiddleware(jwtSecret string) func(http.Handler) http.Handler { + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + bearerToken := ExtractBearerToken(r) + result := ParseAndValidateJWT(bearerToken, auth.TokenTypeAccess, jwtSecret) + + if result.Error != nil { + dto.RespondWithAPIError(w, *result.Error) + return + } - if result.Role != "admin" { - apiErr := dto.ErrAuthAdminRequired - apiErr.Detail = "Admin privileges required" - dto.RespondWithAPIError(w, apiErr) - return - } + if result.Role != "admin" { + apiErr := dto.ErrAuthAdminRequired + apiErr.Detail = "Admin privileges required" + dto.RespondWithAPIError(w, apiErr) + return + } - next.ServeHTTP(w, CreateUserContext(r, result.UserID, result.Role)) - }) + next.ServeHTTP(w, CreateUserContext(r, result.UserID, result.Role)) + }) + } } // UserAuthMiddleware provides authentication for regular user routes. -func UserAuthMiddleware(next http.Handler) http.Handler { - return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - bearerToken := ExtractBearerToken(r) - result := ParseAndValidateJWT(bearerToken, auth.TokenTypeAccess) - - if result.Error != nil { - dto.RespondWithAPIError(w, *result.Error) - return - } +func UserAuthMiddleware(jwtSecret string) func(http.Handler) http.Handler { + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + bearerToken := ExtractBearerToken(r) + result := ParseAndValidateJWT(bearerToken, auth.TokenTypeAccess, jwtSecret) + + if result.Error != nil { + dto.RespondWithAPIError(w, *result.Error) + return + } - next.ServeHTTP(w, CreateUserContext(r, result.UserID, result.Role)) - }) + next.ServeHTTP(w, CreateUserContext(r, result.UserID, result.Role)) + }) + } } // AdminOnlyHandler wraps a handler to require admin role. diff --git a/api/models/models.go b/api/models/models.go index 1e8389eb..d1c16adc 100644 --- a/api/models/models.go +++ b/api/models/models.go @@ -1,7 +1,7 @@ package models import ( - "log" + "log/slog" "github.com/pkoukk/tiktoken-go" ) @@ -29,7 +29,7 @@ func (m Message) TokenCount() int32 { } else { tokenCount, err := getTokenCount(m.Content) if err != nil { - log.Println(err) + slog.Info("token count failed", "error", err) } return int32(tokenCount) + 1 } diff --git a/api/pkg/app/app.go b/api/pkg/app/app.go new file mode 100644 index 00000000..9de53aca --- /dev/null +++ b/api/pkg/app/app.go @@ -0,0 +1,28 @@ +// Package app provides the central application context holding all shared dependencies. +// It replaces package-level global state with dependency injection. +package app + +import ( + "database/sql" + + "github.com/swuecho/chat_backend/sqlc_queries" + "golang.org/x/time/rate" +) + +// App holds all shared application dependencies. +// Pass this to constructors instead of relying on global variables. +type App struct { + DB *sql.DB + Q *sqlc_queries.Queries + Config Config + RateLimiter *rate.Limiter +} + +// Config holds runtime configuration values needed across packages. +type Config struct { + OpenAIKey string + OpenAIProxy string + JWTSecret string + JWTAudience string + DefaultLimit int32 +} diff --git a/api/pkg/util/util.go b/api/pkg/util/util.go new file mode 100644 index 00000000..f0febfbb --- /dev/null +++ b/api/pkg/util/util.go @@ -0,0 +1,130 @@ +// Package util provides shared utility functions used across the application. +package util + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "os" + "strconv" + "strings" + + "github.com/google/uuid" + "github.com/pkoukk/tiktoken-go" + "github.com/swuecho/chat_backend/dto" + "github.com/swuecho/chat_backend/middleware" +) + +// NewUUID generates a new UUID v7 string. +func NewUUID() string { + uuidv7, err := uuid.NewV7() + if err != nil { + return uuid.NewString() + } + return uuidv7.String() +} + +// TokenCount returns the estimated token count for a text string. +func TokenCount(content string) (int, error) { + tke, err := tiktoken.GetEncoding("cl100k_base") + if err != nil { + return 0, err + } + return len(tke.Encode(content, nil, nil)), nil +} + +// FirstNWords returns the first n words of a string. +func FirstNWords(s string, n int) string { + if s == "" { + return "" + } + words := strings.Fields(s) + if len(words) <= n { + return s + } + return strings.Join(words[:n], " ") +} + +// SetupSSE configures the response writer for Server-Sent Events +// and returns a Flusher for streaming. +func SetupSSE(w http.ResponseWriter) (http.Flusher, error) { + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache, no-transform") + w.Header().Set("Connection", "keep-alive") + w.Header().Set("Access-Control-Allow-Origin", "*") + w.Header().Set("X-Accel-Buffering", "no") + w.Header().Del("Content-Length") + w.Header().Del("Content-Encoding") + + flusher, ok := w.(http.Flusher) + if !ok { + return nil, errors.New(dto.ErrorStreamUnsupported) + } + return flusher, nil +} + +// PerWordStreamLimit returns the per-word streaming limit from the environment. +func PerWordStreamLimit() int { + val := os.Getenv("PER_WORD_STREAM_LIMIT") + if val == "" { + return 200 + } + n, err := strconv.Atoi(val) + if err != nil { + return 200 + } + return n +} + +// PaginationParams extracts limit and offset from HTTP query parameters. +func PaginationParams(r *http.Request) (limit int32, offset int32) { + limit = 100 + if v := r.URL.Query().Get("limit"); v != "" { + if l, err := strconv.ParseInt(v, 10, 32); err == nil { + limit = int32(l) + } + } + if v := r.URL.Query().Get("offset"); v != "" { + if o, err := strconv.ParseInt(v, 10, 32); err == nil { + offset = int32(o) + } + } + return +} + +// LimitParam extracts an optional limit parameter with a default value. +func LimitParam(r *http.Request, defaultLimit int32) int32 { + v := r.URL.Query().Get("limit") + if v == "" { + return defaultLimit + } + n, err := strconv.ParseInt(v, 10, 32) + if err != nil { + return defaultLimit + } + return int32(n) +} + +// DecodeJSON decodes a JSON request body into the target value. +func DecodeJSON(r *http.Request, target any) error { + return json.NewDecoder(r.Body).Decode(target) +} + +// UserID extracts the authenticated user ID from the request context. +func UserID(ctx context.Context) (int32, error) { + val := ctx.Value(middleware.UserContextKey) + if val == nil { + return 0, fmt.Errorf("no user ID in context") + } + userIDStr, ok := val.(string) + if !ok { + return 0, fmt.Errorf("user ID in context is not a string") + } + userIDInt, err := strconv.ParseInt(userIDStr, 10, 32) + if err != nil { + return 0, fmt.Errorf("invalid user ID: %s", userIDStr) + } + return int32(userIDInt), nil +} diff --git a/api/provider/llm_openai.go b/api/provider/llm_openai.go index 74b18a04..08e5b36b 100644 --- a/api/provider/llm_openai.go +++ b/api/provider/llm_openai.go @@ -3,7 +3,7 @@ package provider import ( "encoding/base64" "fmt" - "log" + "log/slog" "net/http" "net/url" "os" @@ -69,13 +69,13 @@ func messagesToOpenAIMesages(messages []models.Message, chatFiles []sqlc_queries firstUserMessage, idx, found := lo.FindIndexOf(open_ai_msgs, func(msg openai.ChatCompletionMessage) bool { return msg.Role == "user" }) if found { - log.Printf("firstUserMessage: %+v\n", firstUserMessage) + slog.Info("firstUserMessage before attach", "msg", firstUserMessage) open_ai_msgs[idx].MultiContent = append( []openai.ChatMessagePart{ {Type: openai.ChatMessagePartTypeText, Text: firstUserMessage.Content}, }, parts...) open_ai_msgs[idx].Content = "" - log.Printf("firstUserMessage: %+v\n", firstUserMessage) + slog.Info("firstUserMessage after attach", "msg", firstUserMessage) } return open_ai_msgs @@ -117,7 +117,7 @@ func NormalizeOpenAIModelName(chatModel sqlc_queries.ChatModel, modelName string if strings.Contains(chatModel.Url, "open.bigmodel.cn") { normalized := strings.ToLower(modelName) if normalized != modelName { - log.Printf("Normalizing BigModel model name from %q to %q", modelName, normalized) + slog.Info("Normalizing BigModel model name from %q to %q", modelName, normalized) } return normalized } @@ -129,7 +129,7 @@ func configOpenAIProxy(clientCfg *openai.ClientConfig, proxyURL string) { if proxyUrlStr != "" { proxyUrl, err := url.Parse(proxyUrlStr) if err != nil { - log.Printf("Error parsing proxy URL: %v", err) + slog.Error("error parsing proxy URL", "error", err) } transport := &http.Transport{ Proxy: http.ProxyURL(proxyUrl), @@ -148,7 +148,7 @@ func GenOpenAIConfig(chatModel sqlc_queries.ChatModel, cfg Config) (openai.Clien if err != nil { return openai.ClientConfig{}, err } - log.Printf("OpenAI-compatible URL resolved - Model: %s, ConfiguredURL: %s, BaseURL: %s", chatModel.Name, chatModel.Url, baseUrl) + slog.Info("OpenAI-compatible URL resolved", "model", chatModel.Name, "configuredURL", chatModel.Url, "baseURL", baseUrl) var config openai.ClientConfig if os.Getenv("AZURE_RESOURCE_NAME") != "" { diff --git a/api/provider/llm_summary.go b/api/provider/llm_summary.go index 84205faa..d3797a7b 100644 --- a/api/provider/llm_summary.go +++ b/api/provider/llm_summary.go @@ -2,7 +2,7 @@ package provider import ( "context" - "log" + "log/slog" "strings" "time" @@ -12,29 +12,26 @@ import ( "github.com/tmc/langchaingo/textsplitter" ) -// OpenAIToken is set during initialization to avoid importing main package globals. -var OpenAIToken string - // SummarizeWithTimeout generates a summary with a 20-second timeout. -func SummarizeWithTimeout(baseURL, content string) string { +func SummarizeWithTimeout(apiToken, baseURL, content string) string { // Create a context with a 20 second timeout ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) defer cancel() // Call the summarize function with the context - summary := llm_summarize(ctx, baseURL, content) + summary := llm_summarize(ctx, apiToken, baseURL, content) return summary } -func llm_summarize(ctx context.Context, baseURL string, doc string) string { +func llm_summarize(ctx context.Context, apiToken, baseURL string, doc string) string { baseURL = strings.TrimSuffix(baseURL, "/v1") llm, err := openai.New( - openai.WithToken(OpenAIToken), + openai.WithToken(apiToken), openai.WithBaseURL(baseURL), ) if err != nil { - log.Printf("failed to create openai client %s: %v", baseURL, err) + slog.Info("failed to create openai client %s: %v", baseURL, err) return "" } @@ -44,7 +41,7 @@ func llm_summarize(ctx context.Context, baseURL string, doc string) string { ) outputValues, err := chains.Call(ctx, llmSummarizationChain, map[string]any{"input_documents": docs}) if err != nil { - log.Printf("failed to call chain: %s, %v", baseURL, err) + slog.Info("failed to call chain: %s, %v", baseURL, err) return "" } out, _ := outputValues["text"].(string) diff --git a/api/provider/model_claude3_service.go b/api/provider/model_claude3_service.go index ac50d335..17659c45 100644 --- a/api/provider/model_claude3_service.go +++ b/api/provider/model_claude3_service.go @@ -8,16 +8,15 @@ import ( "errors" "fmt" "io" - "log" + "log/slog" "net/http" "os" "time" - openai "github.com/sashabaranov/go-openai" + "github.com/swuecho/chat_backend/dto" claude "github.com/swuecho/chat_backend/llm/claude" "github.com/swuecho/chat_backend/models" "github.com/swuecho/chat_backend/sqlc_queries" - "github.com/swuecho/chat_backend/dto" ) // ClaudeResponse represents the response structure from Claude API @@ -41,45 +40,31 @@ func NewClaude3ChatModel(h Handler) *Claude3ChatModel { return &Claude3ChatModel{h: h} } -func (m *Claude3ChatModel) Stream(w http.ResponseWriter, chatSession sqlc_queries.ChatSession, chat_compeletion_messages []models.Message, chatUuid string, regenerate bool, stream bool) (*models.LLMAnswer, error) { - // Get chat model configuration - chatModel, err := GetChatModel(m.h.RequestContext(), m.h.Queries(), chatSession.Model) +func (m *Claude3ChatModel) Stream(ctx context.Context, chatSession sqlc_queries.ChatSession, chatCompletionMessages []models.Message, chatUuid string, regenerate bool, stream bool) (<-chan StreamChunk, error) { + chatModel, err := GetChatModel(ctx, m.h.Queries(), chatSession.Model) if err != nil { return nil, err } - // Get chat files if any - chatFiles, err := GetChatFiles(m.h.RequestContext(), m.h.Queries(), chatSession.Uuid) + chatFiles, err := GetChatFiles(ctx, m.h.Queries(), chatSession.Uuid) if err != nil { return nil, err } - // create a new strings.Builder - // iterate through the messages and format them - // print the user's question - // convert assistant's response to json format - // "messages": [ - // {"role": "user", "content": "Hello, world"} - // ] - // first message is user instead of system - var messages []openai.ChatCompletionMessage - if len(chat_compeletion_messages) > 1 { - // first message used as system message - // messages start with second message - // drop the first assistant message if it is an assistant message - claude_messages := chat_compeletion_messages[1:] - - if len(claude_messages) > 0 && claude_messages[0].Role == "assistant" { - claude_messages = claude_messages[1:] + var claudeMessages []models.Message + if len(chatCompletionMessages) > 1 { + claudeMessages = chatCompletionMessages[1:] + if len(claudeMessages) > 0 && claudeMessages[0].Role == "assistant" { + claudeMessages = claudeMessages[1:] } - messages = messagesToOpenAIMesages(claude_messages, chatFiles) } else { - // only system message, return and do nothing return nil, dto.ErrSystemMessageError } - // Prepare request payload + + messages := messagesToOpenAIMesages(claudeMessages, chatFiles) + jsonData := map[string]any{ - "system": chat_compeletion_messages[0].Content, + "system": chatCompletionMessages[0].Content, "model": chatSession.Model, "messages": messages, "max_tokens": chatSession.MaxTokens, @@ -93,18 +78,12 @@ func (m *Claude3ChatModel) Stream(w http.ResponseWriter, chatSession sqlc_querie return nil, dto.ErrValidationInvalidInputGeneric.WithDetail("failed to marshal request payload").WithDebugInfo(err.Error()) } - // Get request context for cancellation support - ctx := m.h.RequestContext() - - // Create HTTP request with context req, err := http.NewRequestWithContext(ctx, "POST", chatModel.Url, bytes.NewBuffer(jsonValue)) if err != nil { return nil, dto.ErrClaudeRequestFailed.WithDetail("failed to create HTTP request").WithDebugInfo(err.Error()) } - // add headers to the request apiKey := os.Getenv(chatModel.ApiAuthKey) - if apiKey == "" { return nil, dto.ErrAuthInvalidCredentials.WithDetail(fmt.Sprintf("missing API key for model %s", chatSession.Model)) } @@ -117,40 +96,38 @@ func (m *Claude3ChatModel) Stream(w http.ResponseWriter, chatSession sqlc_querie req.Header.Set("Content-Type", "application/json") req.Header.Set("anthropic-version", "2023-06-01") - if !stream { - req.Header.Set("Accept", "application/json") - client := http.Client{ - Timeout: 5 * time.Minute, - } - - llmAnswer, err := doGenerateClaude3(ctx, client, req) - if err != nil { - return nil, dto.ErrClaudeRequestFailed.WithDetail("failed to generate response").WithDebugInfo(err.Error()) - } - - answerResponse := buildStreamResponse(llmAnswer.AnswerId, llmAnswer.Answer) - data, err := json.Marshal(answerResponse) - if err != nil { - return nil, dto.ErrInternalUnexpected.WithDetail("failed to marshal response").WithDebugInfo(err.Error()) - } - - if _, err := fmt.Fprint(w, string(data)); err != nil { - return nil, dto.ErrClaudeResponseFailed.WithDetail("failed to write response").WithDebugInfo(err.Error()) - } + ch := make(chan StreamChunk, 10) - return llmAnswer, nil - } - - // Handle streaming response - req.Header.Set("Accept", "text/event-stream") - req.Header.Set("Cache-Control", "no-cache") - req.Header.Set("Connection", "keep-alive") - - llmAnswer, err := chatStreamClaude3(ctx, w, req, chatUuid, regenerate) - if err != nil { - return nil, dto.ErrClaudeStreamFailed.WithDetail("failed to stream response").WithDebugInfo(err.Error()) - } - return llmAnswer, nil + if !stream { + go func() { + defer close(ch) + req.Header.Set("Accept", "application/json") + client := http.Client{Timeout: 5 * time.Minute} + llmAnswer, err := doGenerateClaude3(ctx, client, req) + if err != nil { + ch <- StreamChunk{Err: err} + return + } + ch <- StreamChunk{ + ID: llmAnswer.AnswerId, + Done: true, + FinalAnswer: &models.LLMAnswer{ + Answer: llmAnswer.Answer, + AnswerId: llmAnswer.AnswerId, + }, + } + }() + return ch, nil + } + + go func() { + defer close(ch) + req.Header.Set("Accept", "text/event-stream") + req.Header.Set("Cache-Control", "no-cache") + req.Header.Set("Connection", "keep-alive") + chatStreamClaude3(ctx, ch, req, chatUuid, regenerate) + }() + return ch, nil } func doGenerateClaude3(ctx context.Context, client http.Client, req *http.Request) (*models.LLMAnswer, error) { @@ -158,140 +135,99 @@ func doGenerateClaude3(ctx context.Context, client http.Client, req *http.Reques if err != nil { return nil, dto.ErrClaudeRequestFailed.WithMessage("Failed to process Claude request").WithDebugInfo(err.Error()) } - // Unmarshal directly from resp.Body var message claude.Response if err := json.NewDecoder(resp.Body).Decode(&message); err != nil { + resp.Body.Close() return nil, dto.ErrClaudeInvalidResponse.WithMessage("Failed to unmarshal Claude response").WithDebugInfo(err.Error()) } - defer resp.Body.Close() - uuid := message.ID + resp.Body.Close() firstMessage := message.Content[0].Text return &models.LLMAnswer{ - AnswerId: uuid, + AnswerId: message.ID, Answer: firstMessage, }, nil } -// claude-3-opus-20240229 -// claude-3-sonnet-20240229 -// claude-3-haiku-20240307 -func chatStreamClaude3(ctx context.Context, w http.ResponseWriter, req *http.Request, chatUuid string, regenerate bool) (*models.LLMAnswer, error) { - - // create the http client and send the request - client := &http.Client{ - Timeout: 5 * time.Minute, - } +func chatStreamClaude3(ctx context.Context, ch chan<- StreamChunk, req *http.Request, chatUuid string, regenerate bool) { + client := &http.Client{Timeout: 5 * time.Minute} resp, err := client.Do(req) if err != nil { - return nil, dto.ErrClaudeRequestFailed.WithMessage("Failed to process Claude streaming request").WithDebugInfo(err.Error()) + ch <- StreamChunk{Err: dto.ErrClaudeRequestFailed.WithMessage("Failed to process Claude streaming request").WithDebugInfo(err.Error())} + return } - // Use smaller buffer for more responsive streaming ioreader := bufio.NewReaderSize(resp.Body, 1024) - - // read the response body defer resp.Body.Close() - // loop over the response body and print data - - flusher, err := SetupSSEStream(w) - if err != nil { - return nil, dto.APIError{ - HTTPCode: http.StatusInternalServerError, - Code: "STREAM_UNSUPPORTED", - Message: "Streaming unsupported by client", - } - } - - // Flush immediately to establish connection - flusher.Flush() var answer string - answer_id := generateAnswerID(chatUuid, regenerate) - + answerID := generateAnswerID(chatUuid, regenerate) var headerData = []byte("data: ") count := 0 + for { - // Check if client disconnected or context was cancelled select { case <-ctx.Done(): - log.Printf("Claude stream cancelled by client: %v", ctx.Err()) - // Return current accumulated content when cancelled - return &models.LLMAnswer{Answer: answer, AnswerId: answer_id}, nil + slog.Info("Claude stream cancelled by client", "error", ctx.Err()) + ch <- StreamChunk{Done: true, FinalAnswer: &models.LLMAnswer{Answer: answer, AnswerId: answerID}} + return default: } count++ - // prevent infinite loop if count > 10000 { break } line, err := ioreader.ReadBytes('\n') - log.Printf("%+v", string(line)) if err != nil { if errors.Is(err, io.EOF) { if bytes.HasPrefix(line, []byte("{\"type\":\"error\"")) { - log.Println(string(line)) - err := FlushResponse(w, flusher, StreamingResponse{ - AnswerID: NewUUID(), - Content: string(line), - IsFinal: true, - }) - if err != nil { - log.Printf("Failed to flush error response: %v", err) + ch <- StreamChunk{ + ID: NewUUID(), + Done: true, + FinalAnswer: &models.LLMAnswer{Answer: string(line), AnswerId: answerID}, } } fmt.Println("End of stream reached") - return nil, err + break } - return nil, err + ch <- StreamChunk{Err: err} + return } line = bytes.TrimPrefix(line, headerData) if bytes.HasPrefix(line, []byte("event: message_stop")) { - // stream.isFinished = true - // No need to send full content at the end since we're sending deltas break } if bytes.HasPrefix(line, []byte("{\"type\":\"error\"")) { - log.Println(string(line)) - return nil, dto.ErrClaudeStreamFailed.WithMessage("Error in Claude API response").WithDebugInfo(string(line)) + ch <- StreamChunk{Err: dto.ErrClaudeStreamFailed.WithMessage("Error in Claude API response").WithDebugInfo(string(line))} + return } - if answer_id == "" { - answer_id = NewUUID() + if answerID == "" { + answerID = NewUUID() } if bytes.HasPrefix(line, []byte("{\"type\":\"content_block_start\"")) { - answer = claude.AnswerFromBlockStart(line) - err := FlushResponse(w, flusher, StreamingResponse{ - AnswerID: answer_id, - Content: answer, - IsFinal: false, - }) - if err != nil { - log.Printf("Failed to flush content block start: %v", err) + delta := claude.AnswerFromBlockStart(line) + answer = delta + if len(delta) > 0 { + ch <- StreamChunk{ID: answerID, Content: delta} } } if bytes.HasPrefix(line, []byte("{\"type\":\"content_block_delta\"")) { delta := claude.AnswerFromBlockDelta(line) - answer += delta // Still accumulate for final answer storage - // Send only the delta content - err := FlushResponse(w, flusher, StreamingResponse{ - AnswerID: answer_id, - Content: delta, - IsFinal: false, - }) - if err != nil { - log.Printf("Failed to flush content block delta: %v", err) + answer += delta + if len(delta) > 0 { + ch <- StreamChunk{ID: answerID, Content: delta} } } - // Flush after every iteration to ensure immediate delivery - // This prevents data from being held in buffers - if count%3 == 0 { - flusher.Flush() - } } - return &models.LLMAnswer{ - Answer: answer, - AnswerId: answer_id, - }, nil + + ch <- StreamChunk{ + ID: answerID, + Done: true, + FinalAnswer: &models.LLMAnswer{ + Answer: answer, + AnswerId: answerID, + }, + } } diff --git a/api/provider/model_completion_service.go b/api/provider/model_completion_service.go index 6004e5a0..abc4ca09 100644 --- a/api/provider/model_completion_service.go +++ b/api/provider/model_completion_service.go @@ -1,20 +1,17 @@ package provider import ( - "fmt" "context" - "encoding/json" "errors" "io" - "log" - "net/http" + "log/slog" "strings" openai "github.com/sashabaranov/go-openai" + "github.com/swuecho/chat_backend/dto" "github.com/swuecho/chat_backend/models" "github.com/swuecho/chat_backend/sqlc_queries" - "github.com/swuecho/chat_backend/dto" ) // CompletionChatModel implements ChatModel interface for OpenAI completion models @@ -27,43 +24,38 @@ func NewCompletionChatModel(h Handler) *CompletionChatModel { return &CompletionChatModel{h: h} } -// Stream implements the ChatModel interface for completion model scenarios -func (m *CompletionChatModel) Stream(w http.ResponseWriter, chatSession sqlc_queries.ChatSession, chat_completion_messages []models.Message, chatUuid string, regenerate bool, stream bool) (*models.LLMAnswer, error) { - // Get request context for cancellation support - ctx := m.h.RequestContext() - return m.completionStream(ctx, w, chatSession, chat_completion_messages, chatUuid, regenerate, stream) +func (m *CompletionChatModel) Stream(ctx context.Context, chatSession sqlc_queries.ChatSession, chatCompletionMessages []models.Message, chatUuid string, regenerate bool, stream bool) (<-chan StreamChunk, error) { + ch := make(chan StreamChunk, 10) + go func() { + defer close(ch) + m.completionStream(ctx, ch, chatSession, chatCompletionMessages, chatUuid, regenerate) + }() + return ch, nil } -// completionStream handles streaming for OpenAI completion models -func (m *CompletionChatModel) completionStream(ctx context.Context, w http.ResponseWriter, chatSession sqlc_queries.ChatSession, chat_completion_messages []models.Message, chatUuid string, regenerate bool, _ bool) (*models.LLMAnswer, error) { - // Check per chat_model rate limit - m.h.Config().RateLimiter.Wait(m.h.RequestContext()) +func (m *CompletionChatModel) completionStream(ctx context.Context, ch chan<- StreamChunk, chatSession sqlc_queries.ChatSession, chatCompletionMessages []models.Message, chatUuid string, regenerate bool) { + m.h.Config().RateLimiter.Wait(ctx) - exceedPerModeRateLimitOrError := m.h.CheckModelAccess(w, chatSession.Uuid, chatSession.Model, chatSession.UserID) - if exceedPerModeRateLimitOrError { - return nil, fmt.Errorf("exceed per mode rate limit") + if err := m.h.CheckModelAccess(ctx, chatSession.Uuid, chatSession.Model, chatSession.UserID); err != nil { + ch <- StreamChunk{Err: err} + return } - // Get chat model configuration - chatModel, err := GetChatModel(m.h.RequestContext(), m.h.Queries(), chatSession.Model) + chatModel, err := GetChatModel(ctx, m.h.Queries(), chatSession.Model) if err != nil { - dto.RespondWithAPIError(w, dto.CreateAPIError(dto.ErrResourceNotFound(""), "chat model "+chatSession.Model, "")) - return nil, err + ch <- StreamChunk{Err: err} + return } - // Generate OpenAI client configuration config, err := GenOpenAIConfig(*chatModel, m.h.Config()) if err != nil { - dto.RespondWithAPIError(w, dto.CreateAPIError(dto.ErrInternalUnexpected, "Failed to generate OpenAI configuration", err.Error())) - return nil, err + ch <- StreamChunk{Err: err} + return } client := openai.NewClientWithConfig(config) + prompt := chatCompletionMessages[len(chatCompletionMessages)-1].Content - // Get the latest message content as prompt - prompt := chat_completion_messages[len(chat_completion_messages)-1].Content - - // Create completion request N := chatSession.N req := openai.CompletionRequest{ Model: chatSession.Model, @@ -74,108 +66,71 @@ func (m *CompletionChatModel) completionStream(ctx context.Context, w http.Respo Stream: true, } - // Create completion stream with timeout ctx, cancel := context.WithTimeout(ctx, dto.DefaultRequestTimeout) defer cancel() stream, err := client.CreateCompletionStream(ctx, req) if err != nil { - dto.RespondWithAPIError(w, dto.CreateAPIError(dto.ErrInternalUnexpected, "Failed to create completion stream", err.Error())) - return nil, err + ch <- StreamChunk{Err: dto.ErrInternalUnexpected.WithMessage("Failed to create completion stream").WithDebugInfo(err.Error())} + return } defer stream.Close() - // Setup SSE streaming - flusher, err := SetupSSEStream(w) - if err != nil { - dto.RespondWithAPIError(w, dto.CreateAPIError(dto.ErrInternalUnexpected, "Streaming unsupported by client", err.Error())) - return nil, err - } - var answer string - answer_id := generateAnswerID(chatUuid, regenerate) + answerID := generateAnswerID(chatUuid, regenerate) TextBuffer := NewTextBuffer(N, "```\n"+prompt, "\n```\n") - // Process streaming response for { - // Check if client disconnected or context was cancelled select { case <-ctx.Done(): - log.Printf("Completion stream cancelled by client: %v", ctx.Err()) - // Return current accumulated content when cancelled - return &models.LLMAnswer{Answer: answer, AnswerId: answer_id}, nil + slog.Info("Completion stream cancelled by client", "error", ctx.Err()) + ch <- StreamChunk{Done: true, FinalAnswer: &models.LLMAnswer{Answer: answer, AnswerId: answerID}} + return default: } response, err := stream.Recv() if errors.Is(err, io.EOF) { - // Send the final message if len(answer) > 0 { - err := FlushResponse(w, flusher, StreamingResponse{ - AnswerID: answer_id, - Content: answer, - IsFinal: true, - }) - if err != nil { - log.Printf("Failed to flush final response: %v", err) - } - } - - // Include debug information if enabled - if chatSession.Debug { - req_j, _ := json.Marshal(req) - log.Println(string(req_j)) - answer = answer + "\n" + string(req_j) - err := FlushResponse(w, flusher, StreamingResponse{ - AnswerID: answer_id, - Content: answer, - IsFinal: true, - }) - if err != nil { - log.Printf("Failed to flush debug response: %v", err) - } + ch <- StreamChunk{ID: answerID, Content: answer, Done: true, FinalAnswer: &models.LLMAnswer{AnswerId: answerID, Answer: answer}} + return } break } if err != nil { - dto.RespondWithAPIError(w, dto.ErrChatStreamFailed.WithMessage("Stream error occurred").WithDebugInfo(err.Error())) - return nil, err + ch <- StreamChunk{Err: dto.ErrChatStreamFailed.WithMessage("Stream error occurred").WithDebugInfo(err.Error())} + return } - // Process response chunk textIdx := response.Choices[0].Index delta := response.Choices[0].Text TextBuffer.AppendByIndex(textIdx, delta) if chatSession.Debug { - log.Printf("%d: %s", textIdx, delta) + slog.Info("completion chunk", "index", textIdx, "delta", delta) } - if answer_id == "" { - answer_id = response.ID + if answerID == "" { + answerID = response.ID } - // Concatenate all string builders into a single string answer = TextBuffer.String("\n\n") - // Determine when to flush the response perWordStreamLimit := GetPerWordStreamLimit() if strings.HasSuffix(delta, "\n") || len(answer) < perWordStreamLimit { - if len(answer) == 0 { - log.Print("no content in answer") - } else { - err := FlushResponse(w, flusher, StreamingResponse{ - AnswerID: answer_id, - Content: answer, - IsFinal: false, - }) - if err != nil { - log.Printf("Failed to flush response: %v", err) - } + if len(answer) > 0 { + ch <- StreamChunk{ID: answerID, Content: answer} } } } - return &models.LLMAnswer{AnswerId: answer_id, Answer: answer}, nil + ch <- StreamChunk{ + ID: answerID, + Done: true, + FinalAnswer: &models.LLMAnswer{ + AnswerId: answerID, + Answer: answer, + }, + } } diff --git a/api/provider/model_custom_service.go b/api/provider/model_custom_service.go index 77c1712b..bf1de5f7 100644 --- a/api/provider/model_custom_service.go +++ b/api/provider/model_custom_service.go @@ -8,15 +8,14 @@ import ( "errors" "fmt" "io" - "log" + "log/slog" "net/http" "os" - "strings" + "github.com/swuecho/chat_backend/dto" claude "github.com/swuecho/chat_backend/llm/claude" "github.com/swuecho/chat_backend/models" "github.com/swuecho/chat_backend/sqlc_queries" - "github.com/swuecho/chat_backend/dto" ) // CustomModelResponse represents the response structure for custom models @@ -40,30 +39,27 @@ func NewCustomChatModel(h Handler) *CustomChatModel { return &CustomChatModel{h: h} } -// Stream implements the ChatModel interface for custom model scenarios -func (m *CustomChatModel) Stream(w http.ResponseWriter, chatSession sqlc_queries.ChatSession, chat_completion_messages []models.Message, chatUuid string, regenerate bool, stream bool) (*models.LLMAnswer, error) { - // Get request context for cancellation support - ctx := m.h.RequestContext() - return m.customChatStream(ctx, w, chatSession, chat_completion_messages, chatUuid, regenerate) +func (m *CustomChatModel) Stream(ctx context.Context, chatSession sqlc_queries.ChatSession, chatCompletionMessages []models.Message, chatUuid string, regenerate bool, stream bool) (<-chan StreamChunk, error) { + ch := make(chan StreamChunk, 10) + go func() { + defer close(ch) + m.customChatStream(ctx, ch, chatSession, chatCompletionMessages, chatUuid, regenerate) + }() + return ch, nil } -// customChatStream handles streaming for custom model providers -func (m *CustomChatModel) customChatStream(ctx context.Context, w http.ResponseWriter, chatSession sqlc_queries.ChatSession, chat_completion_messages []models.Message, chatUuid string, regenerate bool) (*models.LLMAnswer, error) { - // Get chat model configuration - chat_model, err := GetChatModel(m.h.RequestContext(), m.h.Queries(), chatSession.Model) +func (m *CustomChatModel) customChatStream(ctx context.Context, ch chan<- StreamChunk, chatSession sqlc_queries.ChatSession, chatCompletionMessages []models.Message, chatUuid string, regenerate bool) { + chatModel, err := GetChatModel(ctx, m.h.Queries(), chatSession.Model) if err != nil { - dto.RespondWithAPIError(w, dto.CreateAPIError(dto.ErrResourceNotFound(""), "chat model: "+chatSession.Model, "")) - return nil, err + ch <- StreamChunk{Err: err} + return } - // Get API key from environment - apiKey := os.Getenv(chat_model.ApiAuthKey) - url := chat_model.Url + apiKey := os.Getenv(chatModel.ApiAuthKey) + url := chatModel.Url - // Format messages for the custom model - prompt := claude.FormatClaudePrompt(chat_completion_messages) + prompt := claude.FormatClaudePrompt(chatCompletionMessages) - // Create request payload jsonData := map[string]any{ "prompt": prompt, "model": chatSession.Model, @@ -73,67 +69,47 @@ func (m *CustomChatModel) customChatStream(ctx context.Context, w http.ResponseW "stream": true, } - // Marshal request data jsonValue, _ := json.Marshal(jsonData) - // Create HTTP request with context req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonValue)) if err != nil { - dto.RespondWithAPIError(w, dto.CreateAPIError(dto.ErrChatRequestFailed, "Failed to create custom model request", err.Error())) - return nil, err + ch <- StreamChunk{Err: dto.ErrChatRequestFailed.WithMessage("Failed to create custom model request").WithDebugInfo(err.Error())} + return } - // Set authentication header if configured - authHeaderName := chat_model.ApiAuthHeader + authHeaderName := chatModel.ApiAuthHeader if authHeaderName != "" { req.Header.Set(authHeaderName, apiKey) } - // Set request headers SetStreamingHeaders(req) - // Send HTTP request client := &http.Client{} resp, err := client.Do(req) if err != nil { - dto.RespondWithAPIError(w, dto.CreateAPIError(dto.ErrChatRequestFailed, "Failed to send custom model request", err.Error())) - return nil, err + ch <- StreamChunk{Err: dto.ErrChatRequestFailed.WithMessage("Failed to send custom model request").WithDebugInfo(err.Error())} + return } defer resp.Body.Close() - // Setup streaming response ioreader := bufio.NewReader(resp.Body) - flusher, err := SetupSSEStream(w) - if err != nil { - dto.RespondWithAPIError(w, dto.CreateAPIError(dto.APIError{ - HTTPCode: http.StatusInternalServerError, - Code: "STREAM_UNSUPPORTED", - Message: "Streaming unsupported by client", - }, "", err.Error())) - return nil, err - } var answer string - var answer_id string + answerID := generateAnswerID(chatUuid, regenerate) var lastFlushLength int - answer_id = generateAnswerID(chatUuid, regenerate) - headerData := []byte("data: ") count := 0 - // Process streaming response for { - // Check if client disconnected or context was cancelled select { case <-ctx.Done(): - log.Printf("Custom model stream cancelled by client: %v", ctx.Err()) - // Return current accumulated content when cancelled - return &models.LLMAnswer{Answer: answer, AnswerId: answer_id}, nil + slog.Info("Custom model stream cancelled by client", "error", ctx.Err()) + ch <- StreamChunk{Done: true, FinalAnswer: &models.LLMAnswer{Answer: answer, AnswerId: answerID}} + return default: } count++ - // Prevent infinite loop if count > 10000 { break } @@ -144,7 +120,8 @@ func (m *CustomChatModel) customChatStream(ctx context.Context, w http.ResponseW fmt.Println("End of stream reached") break } - return nil, err + ch <- StreamChunk{Err: err} + return } if !bytes.HasPrefix(line, headerData) { @@ -154,38 +131,38 @@ func (m *CustomChatModel) customChatStream(ctx context.Context, w http.ResponseW if bytes.HasPrefix(line, []byte("[DONE]")) { fmt.Println("DONE break") - break } - if answer_id == "" { - answer_id = NewUUID() + if answerID == "" { + answerID = NewUUID() } var response CustomModelResponse _ = json.Unmarshal(line, &response) answer = response.Completion - // Determine when to flush the response - shouldFlush := strings.Contains(answer, "\n") || - len(answer) < 200 || - (len(answer)-lastFlushLength) >= 500 - + shouldFlush := len(answer) < 200 || (len(answer)-lastFlushLength) >= 500 if shouldFlush { - err := FlushResponse(w, flusher, StreamingResponse{ - AnswerID: answer_id, - Content: answer, - IsFinal: false, - }) - if err != nil { - log.Printf("Failed to flush response: %v", err) + delta := answer[lastFlushLength:] + if len(delta) > 0 { + ch <- StreamChunk{ID: answerID, Content: delta} } lastFlushLength = len(answer) } } - return &models.LLMAnswer{ - Answer: answer, - AnswerId: answer_id, - }, nil + // Send remaining content + if len(answer) > lastFlushLength { + ch <- StreamChunk{ID: answerID, Content: answer[lastFlushLength:]} + } + + ch <- StreamChunk{ + ID: answerID, + Done: true, + FinalAnswer: &models.LLMAnswer{ + Answer: answer, + AnswerId: answerID, + }, + } } diff --git a/api/provider/model_gemini_service.go b/api/provider/model_gemini_service.go index 9401f030..0b81632a 100644 --- a/api/provider/model_gemini_service.go +++ b/api/provider/model_gemini_service.go @@ -6,30 +6,19 @@ import ( "context" "encoding/json" "errors" - "fmt" "io" - "log" + "log/slog" "net/http" "os" "strings" "time" + "github.com/swuecho/chat_backend/dto" "github.com/swuecho/chat_backend/llm/gemini" "github.com/swuecho/chat_backend/models" "github.com/swuecho/chat_backend/sqlc_queries" - "github.com/swuecho/chat_backend/dto" ) -// Generated by curl-to-Go: https://mholt.github.io/curl-to-go - -// curl https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:generateContent?key=$API_KEY \ -// -H 'Content-Type: application/json' \ -// -X POST \ -// -d '{ -// "contents": [{ -// "parts":[{ -// "text": "Write a story about a magic backpack."}]}]}' 2> /dev/null - // GeminiClient handles communication with the Gemini API type GeminiClient struct { client *http.Client @@ -44,7 +33,7 @@ func NewGeminiClient() *GeminiClient { // Gemini ChatModel implementation type GeminiChatModel struct { - h Handler + h Handler client *GeminiClient } @@ -55,10 +44,10 @@ func NewGeminiChatModel(h Handler) *GeminiChatModel { } } -func (m *GeminiChatModel) Stream(w http.ResponseWriter, chatSession sqlc_queries.ChatSession, messages []models.Message, chatUuid string, regenerate bool, stream bool) (*models.LLMAnswer, error) { +func (m *GeminiChatModel) Stream(ctx context.Context, chatSession sqlc_queries.ChatSession, messages []models.Message, chatUuid string, regenerate bool, stream bool) (<-chan StreamChunk, error) { answerID := generateAnswerID(chatUuid, regenerate) - chatFiles, err := GetChatFiles(m.h.RequestContext(), m.h.Queries(), chatSession.Uuid) + chatFiles, err := GetChatFiles(ctx, m.h.Queries(), chatSession.Uuid) if err != nil { return nil, err } @@ -68,9 +57,6 @@ func (m *GeminiChatModel) Stream(w http.ResponseWriter, chatSession sqlc_queries return nil, dto.ErrInternalUnexpected.WithMessage("Failed to generate Gemini payload").WithDebugInfo(err.Error()) } - // Get request context for cancellation support - ctx := m.h.RequestContext() - url := gemini.BuildAPIURL(chatSession.Model, stream) req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(payloadBytes)) if err != nil { @@ -78,38 +64,49 @@ func (m *GeminiChatModel) Stream(w http.ResponseWriter, chatSession sqlc_queries } req.Header.Set("Content-Type", "application/json") - if stream { - return m.handleStreamResponse(ctx, w, req, answerID) - } + ch := make(chan StreamChunk, 10) - llmAnswer, err := gemini.HandleRegularResponse(*m.client.client, req) - if err != nil { - return nil, dto.ErrInternalUnexpected.WithMessage("Failed to generate regular response").WithDebugInfo(err.Error()) - } - if llmAnswer == nil { - return nil, dto.ErrInternalUnexpected.WithMessage("Empty response from Gemini") + if stream { + go func() { + defer close(ch) + m.handleStreamResponse(ctx, ch, req, answerID) + }() + return ch, nil } - llmAnswer.AnswerId = answerID - response := buildStreamResponse(answerID, llmAnswer.Answer) - data, _ := json.Marshal(response) - fmt.Fprint(w, string(data)) - return llmAnswer, err + go func() { + defer close(ch) + llmAnswer, err := gemini.HandleRegularResponse(*m.client.client, req) + if err != nil { + ch <- StreamChunk{Err: err} + return + } + if llmAnswer == nil { + ch <- StreamChunk{Err: dto.ErrInternalUnexpected.WithMessage("Empty response from Gemini")} + return + } + ch <- StreamChunk{ + ID: answerID, + Done: true, + FinalAnswer: &models.LLMAnswer{ + Answer: llmAnswer.Answer, + AnswerId: answerID, + }, + } + }() + return ch, nil } func GenerateChatTitle(ctx context.Context, model, chatText string) (string, error) { - // Validate API key apiKey := os.Getenv("GEMINI_API_KEY") if apiKey == "" { return "", dto.ErrInternalUnexpected.WithMessage("GEMINI_API_KEY environment variable not set") } - // Validate input if strings.TrimSpace(chatText) == "" { return "", dto.ErrValidationInvalidInput("chat text cannot be empty") } - // Create properly formatted Gemini messages messages := []models.Message{ { Role: "user", @@ -121,13 +118,11 @@ func GenerateChatTitle(ctx context.Context, model, chatText string) (string, err }, } - // Generate proper Gemini payload payloadBytes, err := gemini.GenGemminPayload(messages, nil) if err != nil { return "", dto.ErrInternalUnexpected.WithMessage("Failed to generate Gemini payload").WithDebugInfo(err.Error()) } - // Build URL with proper API key url := gemini.BuildAPIURL(model, false) req, err := http.NewRequest("POST", url, bytes.NewBuffer(payloadBytes)) if err != nil { @@ -139,7 +134,6 @@ func GenerateChatTitle(ctx context.Context, model, chatText string) (string, err return "", dto.ErrInternalUnexpected.WithMessage("Failed to handle Gemini response").WithDebugInfo(err.Error()) } - // Validate and clean up response if answer == nil || answer.Answer == "" { return "", dto.ErrInternalUnexpected.WithMessage("Empty response from Gemini") } @@ -148,13 +142,11 @@ func GenerateChatTitle(ctx context.Context, model, chatText string) (string, err title = strings.Trim(title, `"`) title = strings.Trim(title, `*`) title = strings.Trim(title, `#`) - // Remove common prefixes title = strings.TrimPrefix(title, "Title:") title = strings.TrimPrefix(title, "title:") title = strings.TrimPrefix(title, "Title: ") title = strings.TrimPrefix(title, "title: ") title = strings.TrimSpace(title) - // Remove any remaining markdown or special characters at the start for strings.HasPrefix(title, "#") || strings.HasPrefix(title, "-") || strings.HasPrefix(title, "*") { title = strings.TrimLeft(title, "#-* ") title = strings.TrimSpace(title) @@ -163,65 +155,59 @@ func GenerateChatTitle(ctx context.Context, model, chatText string) (string, err return "", dto.ErrInternalUnexpected.WithMessage("Invalid title generated") } - // Truncate and return return FirstN(title, 100), nil } -func (m *GeminiChatModel) handleStreamResponse(ctx context.Context, w http.ResponseWriter, req *http.Request, answerID string) (*models.LLMAnswer, error) { +func (m *GeminiChatModel) handleStreamResponse(ctx context.Context, ch chan<- StreamChunk, req *http.Request, answerID string) { resp, err := m.client.client.Do(req) if err != nil { - return nil, dto.ErrInternalUnexpected.WithMessage("Failed to send Gemini API request").WithDebugInfo(err.Error()) + ch <- StreamChunk{Err: dto.ErrInternalUnexpected.WithMessage("Failed to send Gemini API request").WithDebugInfo(err.Error())} + return } defer resp.Body.Close() - flusher, err := SetupSSEStream(w) - if err != nil { - return nil, dto.APIError{ - HTTPCode: http.StatusInternalServerError, - Code: "STREAM_UNSUPPORTED", - Message: "Streaming unsupported by client", - } - } - var answer string - log.Println(resp.StatusCode) + slog.Info("gemini response", "statusCode", resp.StatusCode) if resp.StatusCode != http.StatusOK { errorBody, _ := io.ReadAll(resp.Body) - log.Println(string(errorBody)) + slog.Info("gemini error body", "body", string(errorBody)) var apiError gemini.GoogleApiError if json.Unmarshal(errorBody, &apiError) == nil && apiError.Error.Message != "" { - log.Printf("API returned non-200 status: %d %s. Error: %s", resp.StatusCode, http.StatusText(resp.StatusCode), &apiError) + slog.Warn("API returned non-200 status", "statusCode", resp.StatusCode, "statusText", http.StatusText(resp.StatusCode), "error", &apiError) } else { - log.Printf("API returned non-200 status: %d %s. Body: %s", resp.StatusCode, http.StatusText(resp.StatusCode), string(errorBody)) + slog.Warn("API returned non-200 status", "statusCode", resp.StatusCode, "statusText", http.StatusText(resp.StatusCode), "body", string(errorBody)) } - return nil, dto.APIError{ + ch <- StreamChunk{Err: dto.APIError{ HTTPCode: apiError.Error.Code, Code: apiError.Error.Status, Message: apiError.Error.Message, - } + }} + return } ioreader := bufio.NewReader(resp.Body) headerData := []byte("data: ") for count := 0; count < 10000; count++ { - // Check if client disconnected or context was cancelled select { case <-ctx.Done(): - log.Printf("Gemini stream cancelled by client: %v", ctx.Err()) - // Return current accumulated content when cancelled - return &models.LLMAnswer{Answer: answer, AnswerId: answerID}, nil + slog.Info("Gemini stream cancelled by client", "error", ctx.Err()) + ch <- StreamChunk{Done: true, FinalAnswer: &models.LLMAnswer{Answer: answer, AnswerId: answerID}} + return default: } line, err := ioreader.ReadBytes('\n') if err != nil { if errors.Is(err, io.EOF) { - return &models.LLMAnswer{ - Answer: answer, - AnswerId: answerID, - }, nil + ch <- StreamChunk{ + ID: answerID, + Done: true, + FinalAnswer: &models.LLMAnswer{Answer: answer, AnswerId: answerID}, + } + return } - return nil, dto.ErrInternalUnexpected.WithMessage("Error reading stream").WithDebugInfo(err.Error()) + ch <- StreamChunk{Err: dto.ErrInternalUnexpected.WithMessage("Error reading stream").WithDebugInfo(err.Error())} + return } if !bytes.HasPrefix(line, headerData) { @@ -231,23 +217,19 @@ func (m *GeminiChatModel) handleStreamResponse(ctx context.Context, w http.Respo line = bytes.TrimPrefix(line, headerData) if len(line) > 0 { delta := gemini.ParseRespLineDelta(line) - answer += delta // Accumulate delta for final answer storage - // Send only the delta content + answer += delta if len(delta) > 0 { - err := FlushResponse(w, flusher, StreamingResponse{ - AnswerID: answerID, - Content: delta, - IsFinal: false, - }) - if err != nil { - log.Printf("Failed to flush response: %v", err) - } + ch <- StreamChunk{ID: answerID, Content: delta} } } } - return &models.LLMAnswer{ - AnswerId: answerID, - Answer: answer, - }, nil + ch <- StreamChunk{ + ID: answerID, + Done: true, + FinalAnswer: &models.LLMAnswer{ + AnswerId: answerID, + Answer: answer, + }, + } } diff --git a/api/provider/model_ollama_service.go b/api/provider/model_ollama_service.go index c3c42fd6..5aec3b2c 100644 --- a/api/provider/model_ollama_service.go +++ b/api/provider/model_ollama_service.go @@ -8,15 +8,15 @@ import ( "errors" "fmt" "io" - "log" + "log/slog" "net/http" "os" "strings" "time" + "github.com/swuecho/chat_backend/dto" "github.com/swuecho/chat_backend/models" "github.com/swuecho/chat_backend/sqlc_queries" - "github.com/swuecho/chat_backend/dto" ) // OllamaResponse represents the response structure from Ollama API @@ -43,34 +43,34 @@ func NewOllamaChatModel(h Handler) *OllamaChatModel { return &OllamaChatModel{h: h} } -func (m *OllamaChatModel) Stream(w http.ResponseWriter, chatSession sqlc_queries.ChatSession, chat_compeletion_messages []models.Message, chatUuid string, regenerate bool, stream bool) (*models.LLMAnswer, error) { - // Get request context for cancellation support - ctx := m.h.RequestContext() - return chatOllamStream(m.h, ctx, w, chatSession, chat_compeletion_messages, chatUuid, regenerate) +func (m *OllamaChatModel) Stream(ctx context.Context, chatSession sqlc_queries.ChatSession, chatCompletionMessages []models.Message, chatUuid string, regenerate bool, stream bool) (<-chan StreamChunk, error) { + ch := make(chan StreamChunk, 10) + go func() { + defer close(ch) + chatOllamStream(ctx, ch, m.h, chatSession, chatCompletionMessages, chatUuid, regenerate) + }() + return ch, nil } -func chatOllamStream(h Handler, ctx context.Context, w http.ResponseWriter, chatSession sqlc_queries.ChatSession, chat_compeletion_messages []models.Message, chatUuid string, regenerate bool) (*models.LLMAnswer, error) { - // set the api key +func chatOllamStream(ctx context.Context, ch chan<- StreamChunk, h Handler, chatSession sqlc_queries.ChatSession, chatCompletionMessages []models.Message, chatUuid string, regenerate bool) { chatModel, err := GetChatModel(ctx, h.Queries(), chatSession.Model) if err != nil { - dto.RespondWithAPIError(w, dto.CreateAPIError(dto.ErrResourceNotFound(""), "chat model: "+chatSession.Model, "")) - return nil, err + ch <- StreamChunk{Err: dto.ErrResourceNotFound("chat model: " + chatSession.Model)} + return } + jsonData := map[string]any{ "model": strings.Replace(chatSession.Model, "ollama-", "", 1), - "messages": chat_compeletion_messages, + "messages": chatCompletionMessages, } - // convert data to json format jsonValue, _ := json.Marshal(jsonData) - // create the request with context - req, err := http.NewRequestWithContext(ctx, "POST", chatModel.Url, bytes.NewBuffer(jsonValue)) + req, err := http.NewRequestWithContext(ctx, "POST", chatModel.Url, bytes.NewBuffer(jsonValue)) if err != nil { - dto.RespondWithAPIError(w, dto.ErrInternalUnexpected.WithMessage("Failed to make request").WithDebugInfo(err.Error())) - return nil, err + ch <- StreamChunk{Err: dto.ErrInternalUnexpected.WithMessage("Failed to make request").WithDebugInfo(err.Error())} + return } - // add headers to the request apiKey := os.Getenv(chatModel.ApiAuthKey) authHeaderName := chatModel.ApiAuthHeader if authHeaderName != "" { @@ -78,55 +78,35 @@ func chatOllamStream(h Handler, ctx context.Context, w http.ResponseWriter, chat } req.Header.Set("Content-Type", "application/json") - - // set the streaming flag req.Header.Set("Accept", "text/event-stream") req.Header.Set("Cache-Control", "no-cache") req.Header.Set("Connection", "keep-alive") req.Header.Set("Access-Control-Allow-Origin", "*") - // create the http client and send the request - client := &http.Client{ - Timeout: 5 * time.Minute, - } + client := &http.Client{Timeout: 5 * time.Minute} resp, err := client.Do(req) if err != nil { - dto.RespondWithAPIError(w, dto.ErrInternalUnexpected.WithMessage("Failed to create chat completion stream").WithDebugInfo(err.Error())) - return nil, err + ch <- StreamChunk{Err: dto.ErrInternalUnexpected.WithMessage("Failed to create chat completion stream").WithDebugInfo(err.Error())} + return } ioreader := bufio.NewReader(resp.Body) - - // read the response body defer resp.Body.Close() - // loop over the response body and print data - - flusher, err := SetupSSEStream(w) - if err != nil { - dto.RespondWithAPIError(w, dto.APIError{ - HTTPCode: http.StatusInternalServerError, - Code: "STREAM_UNSUPPORTED", - Message: "Streaming unsupported by client", - }) - return nil, err - } var answer string - answer_id := generateAnswerID(chatUuid, regenerate) + answerID := generateAnswerID(chatUuid, regenerate) count := 0 for { - // Check if client disconnected or context was cancelled select { case <-ctx.Done(): - log.Printf("Ollama stream cancelled by client: %v", ctx.Err()) - // Return current accumulated content when cancelled - return &models.LLMAnswer{Answer: answer, AnswerId: answer_id}, nil + slog.Info("Ollama stream cancelled by client", "error", ctx.Err()) + ch <- StreamChunk{Done: true, FinalAnswer: &models.LLMAnswer{Answer: answer, AnswerId: answerID}} + return default: } count++ - // prevent infinite loop if count > 10000 { break } @@ -134,43 +114,38 @@ func chatOllamStream(h Handler, ctx context.Context, w http.ResponseWriter, chat if err != nil { if errors.Is(err, io.EOF) { fmt.Println("End of stream reached") - break // Exit loop if end of stream + break } - return nil, err + ch <- StreamChunk{Err: err} + return } var streamResp OllamaResponse - err = json.Unmarshal(line, &streamResp) - if err != nil { - return nil, err + if err := json.Unmarshal(line, &streamResp); err != nil { + ch <- StreamChunk{Err: err} + return } delta := strings.ReplaceAll(streamResp.Message.Content, "<0x0A>", "\n") - answer += delta // Still accumulate for final answer storage + answer += delta if streamResp.Done { - // stream.isFinished = true fmt.Println("DONE break") - // No need to send full content at the end since we're sending deltas break } - if answer_id == "" { - answer_id = NewUUID() + if answerID == "" { + answerID = NewUUID() } - // Send delta content immediately when available if len(delta) > 0 { - err := FlushResponse(w, flusher, StreamingResponse{ - AnswerID: answer_id, - Content: delta, - IsFinal: false, - }) - if err != nil { - log.Printf("Failed to flush response: %v", err) - } + ch <- StreamChunk{ID: answerID, Content: delta} } } - return &models.LLMAnswer{ - Answer: answer, - AnswerId: answer_id, - }, nil + ch <- StreamChunk{ + ID: answerID, + Done: true, + FinalAnswer: &models.LLMAnswer{ + Answer: answer, + AnswerId: answerID, + }, + } } diff --git a/api/provider/model_openai_service.go b/api/provider/model_openai_service.go index ce699e68..eaaa3879 100644 --- a/api/provider/model_openai_service.go +++ b/api/provider/model_openai_service.go @@ -6,15 +6,14 @@ import ( "errors" "fmt" "io" - "log" - "net/http" + "log/slog" "strings" "time" openai "github.com/sashabaranov/go-openai" + "github.com/swuecho/chat_backend/dto" llm_openai "github.com/swuecho/chat_backend/llm/openai" "github.com/swuecho/chat_backend/models" - "github.com/swuecho/chat_backend/dto" "github.com/swuecho/chat_backend/sqlc_queries" ) @@ -28,27 +27,24 @@ func NewOpenAIChatModel(h Handler) *OpenAIChatModel { return &OpenAIChatModel{h: h} } -func (m *OpenAIChatModel) Stream(w http.ResponseWriter, chatSession sqlc_queries.ChatSession, chatCompletionMessages []models.Message, chatUuid string, regenerate bool, streamOutput bool) (*models.LLMAnswer, error) { - m.h.Config().RateLimiter.Wait(m.h.RequestContext()) +func (m *OpenAIChatModel) Stream(ctx context.Context, chatSession sqlc_queries.ChatSession, chatCompletionMessages []models.Message, chatUuid string, regenerate bool, streamOutput bool) (<-chan StreamChunk, error) { + m.h.Config().RateLimiter.Wait(ctx) - exceedPerModeRateLimitOrError := m.h.CheckModelAccess(w, chatSession.Uuid, chatSession.Model, chatSession.UserID) - if exceedPerModeRateLimitOrError { - return nil, fmt.Errorf("exceed per mode rate limit") + if err := m.h.CheckModelAccess(ctx, chatSession.Uuid, chatSession.Model, chatSession.UserID); err != nil { + return nil, err } - chatModel, err := GetChatModel(m.h.RequestContext(), m.h.Queries(), chatSession.Model) + chatModel, err := GetChatModel(ctx, m.h.Queries(), chatSession.Model) if err != nil { return nil, err } config, err := GenOpenAIConfig(*chatModel, m.h.Config()) - log.Printf("%+v", config.String()) - // print all config details if err != nil { return nil, dto.ErrOpenAIConfigFailed.WithMessage("Failed to generate OpenAI config").WithDebugInfo(err.Error()) } - chatFiles, err := GetChatFiles(m.h.RequestContext(), m.h.Queries(), chatSession.Uuid) + chatFiles, err := GetChatFiles(ctx, m.h.Queries(), chatSession.Uuid) if err != nil { return nil, err } @@ -58,152 +54,134 @@ func (m *OpenAIChatModel) Stream(w http.ResponseWriter, chatSession sqlc_queries if len(openaiReq.Messages) <= 1 { return nil, dto.ErrSystemMessageError } - log.Printf("OpenAI request prepared - Model: %s, MessageCount: %d, Temperature: %.2f", - openaiReq.Model, len(openaiReq.Messages), openaiReq.Temperature) + slog.Info("OpenAI request prepared", "model", openaiReq.Model, "messageCount", len(openaiReq.Messages), "temperature", openaiReq.Temperature) client := openai.NewClientWithConfig(config) - if streamOutput { - return doChatStream(w, client, openaiReq, chatSession.N, chatUuid, regenerate, m.h, chatModel.Url, config.BaseURL) - } else { - return handleRegularResponse(m.h.RequestContext(), w, client, openaiReq, chatModel.Url, config.BaseURL) - } + ch := make(chan StreamChunk, 10) + go func() { + defer close(ch) + if streamOutput { + doChatStream(ctx, ch, client, openaiReq, chatSession.N, chatUuid, regenerate, chatModel.Url, config.BaseURL) + } else { + handleRegularResponse(ctx, ch, client, openaiReq, chatModel.Url, config.BaseURL) + } + }() + return ch, nil } -func handleRegularResponse(ctx context.Context, w http.ResponseWriter, client *openai.Client, req openai.ChatCompletionRequest, configuredURL string, baseURL string) (*models.LLMAnswer, error) { - // check per chat_model limit +func handleRegularResponse(ctx context.Context, ch chan<- StreamChunk, client *openai.Client, req openai.ChatCompletionRequest, configuredURL string, baseURL string) { ctx, cancel := context.WithTimeout(ctx, 5*time.Minute) defer cancel() completion, err := client.CreateChatCompletion(ctx, req) if err != nil { - log.Printf("OpenAI request failed - Model: %s, ConfiguredURL: %s, BaseURL: %s, Error: %+v", req.Model, configuredURL, baseURL, err) - return nil, dto.ErrOpenAIRequestFailed.WithMessage("Failed to create chat completion").WithDebugInfo(err.Error()) + slog.Info("OpenAI request failed - Model: %s, ConfiguredURL: %s, BaseURL: %s, Error: %+v", req.Model, configuredURL, baseURL, err) + ch <- StreamChunk{Err: dto.ErrOpenAIRequestFailed.WithMessage("Failed to create chat completion").WithDebugInfo(err.Error())} + return + } + + ch <- StreamChunk{ + ID: completion.ID, + Content: completion.Choices[0].Message.Content, + Done: true, + FinalAnswer: &models.LLMAnswer{ + Answer: completion.Choices[0].Message.Content, + AnswerId: completion.ID, + }, } - log.Printf("completion: %+v", completion) - data, _ := json.Marshal(completion) - fmt.Fprint(w, string(data)) - return &models.LLMAnswer{Answer: completion.Choices[0].Message.Content, AnswerId: completion.ID}, nil } -// doChatStream handles streaming chat completion responses from OpenAI -// It properly manages thinking tags for models that support reasoning content -func doChatStream(w http.ResponseWriter, client *openai.Client, req openai.ChatCompletionRequest, bufferLen int32, chatUuid string, regenerate bool, handler Handler, configuredURL string, baseURL string) (*models.LLMAnswer, error) { - // Use request context with timeout - ctx, cancel := context.WithTimeout(handler.RequestContext(), 5*time.Minute) +// doChatStream handles streaming chat completion responses from OpenAI. +// It sends chunks on the provided channel and closes it when done. +func doChatStream(ctx context.Context, ch chan<- StreamChunk, client *openai.Client, req openai.ChatCompletionRequest, bufferLen int32, chatUuid string, regenerate bool, configuredURL string, baseURL string) { + ctx, cancel := context.WithTimeout(ctx, 5*time.Minute) defer cancel() - log.Print("Creating OpenAI stream") + slog.Info("Creating OpenAI stream") stream, err := client.CreateChatCompletionStream(ctx, req) - if err != nil { - log.Printf("OpenAI stream setup failed - Model: %s, ConfiguredURL: %s, BaseURL: %s, Error: %+v", req.Model, configuredURL, baseURL, err) - return nil, dto.ErrOpenAIStreamFailed.WithMessage("Failed to create chat completion stream").WithDebugInfo(err.Error()) + slog.Info("OpenAI stream setup failed - Model: %s, ConfiguredURL: %s, BaseURL: %s, Error: %+v", req.Model, configuredURL, baseURL, err) + ch <- StreamChunk{Err: dto.ErrOpenAIStreamFailed.WithMessage("Failed to create chat completion stream").WithDebugInfo(err.Error())} + return } defer func() { if err := stream.Close(); err != nil { - log.Printf("Error closing OpenAI stream: %v", err) + slog.Error("error closing OpenAI stream", "error", err) } }() - // Setup Server-Sent Events (SSE) streaming - flusher, err := SetupSSEStream(w) - if err != nil { - return nil, dto.APIError{ - HTTPCode: http.StatusInternalServerError, - Code: "STREAM_UNSUPPORTED", - Message: "Streaming unsupported by client", - } - } - - // Initialize streaming state - var answer_id string + var answerID string + var hasReason bool + var reasonTagOpened bool + var reasonTagClosed bool - var hasReason bool // Whether we've detected any reasoning content - var reasonTagOpened bool // Whether we've sent the opening tag - var reasonTagClosed bool // Whether we've sent the closing tag - - // Ensure minimum buffer length if bufferLen == 0 { - log.Println("Buffer length is 0, setting to 1") + slog.Info("Buffer length is 0, setting to 1") bufferLen = 1 } - // Initialize buffers for accumulating content TextBuffer := NewTextBuffer(bufferLen, "", "") reasonBuffer := NewTextBuffer(bufferLen, "\n\n", "\n\n\n\n") - answer_id = generateAnswerID(chatUuid, regenerate) - // Main streaming loop + answerID = generateAnswerID(chatUuid, regenerate) + for { - // Check if client disconnected or context was cancelled select { case <-ctx.Done(): - log.Printf("Stream cancelled by client: %v", ctx.Err()) - // Return current accumulated content when cancelled - llmAnswer := models.LLMAnswer{Answer: TextBuffer.String("\n"), AnswerId: answer_id} + slog.Info("Stream cancelled by client", "error", ctx.Err()) + llmAnswer := models.LLMAnswer{Answer: TextBuffer.String("\n"), AnswerId: answerID} if hasReason { llmAnswer.ReasoningContent = reasonBuffer.String("\n") } - return &llmAnswer, nil + ch <- StreamChunk{Done: true, FinalAnswer: &llmAnswer} + return default: } rawLine, err := stream.RecvRaw() if err != nil { - log.Printf("OpenAI stream receive error - Model: %s, ConfiguredURL: %s, BaseURL: %s, Error: %+v", req.Model, configuredURL, baseURL, err) + slog.Info("OpenAI stream receive error - Model: %s, ConfiguredURL: %s, BaseURL: %s, Error: %+v", req.Model, configuredURL, baseURL, err) if errors.Is(err, io.EOF) { if TextBuffer.String("\n") == "" && reasonBuffer.String("\n") == "" { errMsg := fmt.Sprintf("stream closed without content; verify configured URL %q resolves to a valid OpenAI-compatible base URL %q and that model %q is valid", configuredURL, baseURL, req.Model) - log.Print(errMsg) - return nil, dto.ErrOpenAIStreamFailed.WithMessage("Stream closed without content").WithDebugInfo(errMsg) + slog.Info(errMsg) + ch <- StreamChunk{Err: dto.ErrOpenAIStreamFailed.WithMessage("Stream closed without content").WithDebugInfo(errMsg)} + return } - // Stream ended successfully - return accumulated content - llmAnswer := models.LLMAnswer{Answer: TextBuffer.String("\n"), AnswerId: answer_id} + llmAnswer := models.LLMAnswer{Answer: TextBuffer.String("\n"), AnswerId: answerID} if hasReason { llmAnswer.ReasoningContent = reasonBuffer.String("\n") } - return &llmAnswer, nil - } else { - log.Printf("Stream error: %v", err) - return nil, dto.ErrOpenAIStreamFailed.WithMessage("Stream error occurred").WithDebugInfo(err.Error()) + ch <- StreamChunk{Done: true, FinalAnswer: &llmAnswer} + return } + slog.Info("Stream error", "error", err) + ch <- StreamChunk{Err: dto.ErrOpenAIStreamFailed.WithMessage("Stream error occurred").WithDebugInfo(err.Error())} + return } - // Parse the streaming response + response := llm_openai.ChatCompletionStreamResponse{} - err = json.Unmarshal(rawLine, &response) - if err != nil { - log.Printf("Could not unmarshal response: %v\n", err) + if err := json.Unmarshal(rawLine, &response); err != nil { + slog.Info("Could not unmarshal response", "error", err) continue } - // Extract delta content from the response textIdx := response.Choices[0].Index delta := response.Choices[0].Delta - // Accumulate content in buffers (for final answer construction) TextBuffer.AppendByIndex(textIdx, delta.Content) if len(delta.ReasoningContent) > 0 { hasReason = true reasonBuffer.AppendByIndex(textIdx, delta.ReasoningContent) } - // Set answer ID from response if not already set - if answer_id == "" { - answer_id = strings.TrimPrefix(response.ID, "chatcmpl-") + if answerID == "" { + answerID = strings.TrimPrefix(response.ID, "chatcmpl-") } - // Process and send delta content if len(delta.Content) > 0 || len(delta.ReasoningContent) > 0 { deltaToSend := processDelta(delta, &reasonTagOpened, &reasonTagClosed, hasReason) if len(deltaToSend) > 0 { - log.Printf("delta: %s", deltaToSend) - err := FlushResponse(w, flusher, StreamingResponse{ - AnswerID: answer_id, - Content: deltaToSend, - IsFinal: false, - }) - if err != nil { - log.Printf("Failed to flush response: %v", err) - } + ch <- StreamChunk{ID: answerID, Content: deltaToSend} } } } @@ -214,22 +192,16 @@ func processDelta(delta llm_openai.ChatCompletionStreamChoiceDelta, reasonTagOpe var deltaToSend string if len(delta.ReasoningContent) > 0 { - // Handle reasoning content if !*reasonTagOpened { - // First time seeing reasoning content, add opening tag deltaToSend = "" + delta.ReasoningContent *reasonTagOpened = true } else { - // Continue reasoning content deltaToSend = delta.ReasoningContent } } else if hasReason && !*reasonTagClosed { - // We had reasoning content before and now we have regular content for the first time - // Close the think tag first, then send the content deltaToSend = "" + delta.Content *reasonTagClosed = true } else { - // Regular content without reasoning deltaToSend = delta.Content } @@ -247,14 +219,13 @@ func NewChatCompletionRequest(chatSession sqlc_queries.ChatSession, chatCompleti for _, m := range openaiMessages { b, _ := m.MarshalJSON() - log.Printf("messages: %+v\n", string(b)) + slog.Info("openai message", "msg", string(b)) } - log.Printf("messages: %+v\n", openaiMessages) - // Ensure TopP is always greater than 0 to prevent API validation errors + slog.Info("openai messages", "messages", openaiMessages) topP := float32(chatSession.TopP) - 0.01 if topP <= 0 { - topP = 0.01 // Minimum valid value + topP = 0.01 } openaiReq := openai.ChatCompletionRequest{ Model: chatSession.Model, diff --git a/api/provider/provider.go b/api/provider/provider.go index 439de974..8a54ddf7 100644 --- a/api/provider/provider.go +++ b/api/provider/provider.go @@ -9,7 +9,6 @@ package provider import ( "context" - "net/http" "golang.org/x/time/rate" @@ -17,25 +16,35 @@ import ( "github.com/swuecho/chat_backend/sqlc_queries" ) +// StreamChunk represents a single chunk in a streaming LLM response. +type StreamChunk struct { + ID string // answer ID + Content string // delta text content + Done bool // true for the terminal chunk + FinalAnswer *models.LLMAnswer // set on Done (nil on error) + Err error // non-nil if a stream error occurred +} + // ChatModel is the interface all LLM providers must implement. +// Stream returns a channel of StreamChunk and an optional immediate error. +// The channel is closed when streaming completes or fails. type ChatModel interface { - Stream(w http.ResponseWriter, session sqlc_queries.ChatSession, + Stream(ctx context.Context, session sqlc_queries.ChatSession, messages []models.Message, chatUuid string, - regenerate bool, stream bool) (*models.LLMAnswer, error) + regenerate bool, stream bool) (<-chan StreamChunk, error) } // Config holds global configuration needed by providers. type Config struct { - OpenAIKey string - OpenAIProxy string - RateLimiter *rate.Limiter - DefaultLimit int + OpenAIKey string + OpenAIProxy string + RateLimiter *rate.Limiter + DefaultLimit int } // Handler provides request-scoped dependencies that providers need. type Handler interface { - RequestContext() context.Context Queries() *sqlc_queries.Queries - CheckModelAccess(w http.ResponseWriter, chatSessionUuid, model string, userID int32) bool + CheckModelAccess(ctx context.Context, chatSessionUuid, model string, userID int32) error Config() Config } diff --git a/api/provider/streaming.go b/api/provider/streaming.go index e124707b..105d2b67 100644 --- a/api/provider/streaming.go +++ b/api/provider/streaming.go @@ -4,18 +4,14 @@ package provider import ( "context" "encoding/json" - "errors" "fmt" "net/http" - "os" - "strconv" "strings" - "github.com/google/uuid" - "github.com/pkoukk/tiktoken-go" openai "github.com/sashabaranov/go-openai" "github.com/swuecho/chat_backend/dto" + "github.com/swuecho/chat_backend/pkg/util" "github.com/swuecho/chat_backend/sqlc_queries" ) @@ -49,21 +45,8 @@ func FlushResponse(w http.ResponseWriter, flusher http.Flusher, response Streami } // SetupSSEStream configures the response writer for Server-Sent Events. -func SetupSSEStream(w http.ResponseWriter) (http.Flusher, error) { - w.Header().Set("Content-Type", "text/event-stream") - w.Header().Set("Cache-Control", "no-cache, no-transform") - w.Header().Set("Connection", "keep-alive") - w.Header().Set("Access-Control-Allow-Origin", "*") - w.Header().Set("X-Accel-Buffering", "no") - w.Header().Del("Content-Length") - w.Header().Del("Content-Encoding") - - flusher, ok := w.(http.Flusher) - if !ok { - return nil, errors.New(dto.ErrorStreamUnsupported) - } - return flusher, nil -} +// Delegates to pkg/util.SetupSSE. +var SetupSSEStream = util.SetupSSE // --- Text buffer for streaming --- @@ -169,13 +152,7 @@ func FirstN(s string, n int) string { // --- Utility functions --- // NewUUID generates a new UUID v7 string. -func NewUUID() string { - uuidv7, err := uuid.NewV7() - if err != nil { - return uuid.NewString() - } - return uuidv7.String() -} +var NewUUID = util.NewUUID // generateAnswerID creates an answer ID or reuses chatUuid in regenerate mode. func generateAnswerID(chatUuid string, regenerate bool) string { @@ -186,23 +163,7 @@ func generateAnswerID(chatUuid string, regenerate bool) string { } // GetTokenCount returns the number of tokens in the given content. -func GetTokenCount(content string) (int, error) { - tke, err := tiktoken.GetEncoding("cl100k_base") - if err != nil { - return 0, err - } - return len(tke.Encode(content, nil, nil)), nil -} +var GetTokenCount = util.TokenCount // GetPerWordStreamLimit returns the per-word stream limit from env or default. -func GetPerWordStreamLimit() int { - s := os.Getenv("PER_WORD_STREAM_LIMIT") - if s == "" { - return 200 - } - n, err := strconv.Atoi(s) - if err != nil { - return 200 - } - return n -} +var GetPerWordStreamLimit = util.PerWordStreamLimit diff --git a/api/provider/test.go b/api/provider/test.go index 78d4916c..c421f6b7 100644 --- a/api/provider/test.go +++ b/api/provider/test.go @@ -1,11 +1,9 @@ package provider import ( + "context" "encoding/json" - "log" - "net/http" - "github.com/swuecho/chat_backend/dto" "github.com/swuecho/chat_backend/models" "github.com/swuecho/chat_backend/sqlc_queries" ) @@ -20,36 +18,33 @@ func NewTestChatModel(h Handler) *TestChatModel { return &TestChatModel{h: h} } -func (m *TestChatModel) Stream(w http.ResponseWriter, session sqlc_queries.ChatSession, - messages []models.Message, chatUuid string, regenerate bool, stream bool) (*models.LLMAnswer, error) { +func (m *TestChatModel) Stream(ctx context.Context, session sqlc_queries.ChatSession, + messages []models.Message, chatUuid string, regenerate bool, stream bool) (<-chan StreamChunk, error) { - chatFiles, err := GetChatFiles(m.h.RequestContext(), m.h.Queries(), session.Uuid) + chatFiles, err := GetChatFiles(ctx, m.h.Queries(), session.Uuid) if err != nil { - dto.RespondWithAPIError(w, dto.ErrInternalUnexpected.WithDetail("Failed to get chat files").WithDebugInfo(err.Error())) return nil, err } answerID := generateAnswerID(chatUuid, regenerate) - flusher, err := SetupSSEStream(w) - if err != nil { - dto.RespondWithAPIError(w, dto.APIError{ - HTTPCode: http.StatusInternalServerError, - Code: "STREAM_UNSUPPORTED", - Message: "Streaming unsupported by client", - }) - return nil, err - } - answer := "Hi, I am a chatbot. I can help you to find the best answer for your question. Please ask me a question." - if err := FlushResponse(w, flusher, StreamingResponse{AnswerID: answerID, Content: answer}); err != nil { - log.Printf("Failed to flush response: %v", err) - } - - if session.Debug { - openaiReq := NewChatCompletionRequest(session, messages, chatFiles, false) - reqJ, _ := json.Marshal(openaiReq) - FlushResponse(w, flusher, StreamingResponse{AnswerID: answerID, Content: answer + "\n" + string(reqJ), IsFinal: true}) - } - return &models.LLMAnswer{Answer: answer, AnswerId: answerID}, nil + ch := make(chan StreamChunk, 2) + go func() { + defer close(ch) + ch <- StreamChunk{ID: answerID, Content: answer} + + if session.Debug { + openaiReq := NewChatCompletionRequest(session, messages, chatFiles, false) + reqJ, _ := json.Marshal(openaiReq) + ch <- StreamChunk{ID: answerID, Content: "\n" + string(reqJ)} + } + + ch <- StreamChunk{ + ID: answerID, + Done: true, + FinalAnswer: &models.LLMAnswer{Answer: answer, AnswerId: answerID}, + } + }() + return ch, nil } diff --git a/api/sqlc_queries/chat_session.sql.go b/api/sqlc_queries/chat_session.sql.go index a2a1e8dd..cb270392 100644 --- a/api/sqlc_queries/chat_session.sql.go +++ b/api/sqlc_queries/chat_session.sql.go @@ -181,20 +181,20 @@ returning id, user_id, uuid, topic, created_at, updated_at, active, model, max_l ` type CreateOrUpdateChatSessionByUUIDParams struct { - Uuid string `json:"uuid"` - UserID int32 `json:"userId"` - Topic string `json:"topic"` - MaxLength int32 `json:"maxLength"` - Temperature float64 `json:"temperature"` - Model string `json:"model"` - MaxTokens int32 `json:"maxTokens"` - TopP float64 `json:"topP"` - N int32 `json:"n"` - Debug bool `json:"debug"` - SummarizeMode bool `json:"summarizeMode"` - WorkspaceID sql.NullInt32 `json:"workspaceId"` - ExploreMode bool `json:"exploreMode"` - ArtifactEnabled bool `json:"artifactEnabled"` + Uuid string `json:"uuid"` + UserID int32 `json:"userId"` + Topic string `json:"topic"` + MaxLength int32 `json:"maxLength"` + Temperature float64 `json:"temperature"` + Model string `json:"model"` + MaxTokens int32 `json:"maxTokens"` + TopP float64 `json:"topP"` + N int32 `json:"n"` + Debug bool `json:"debug"` + SummarizeMode bool `json:"summarizeMode"` + WorkspaceID sql.NullInt32 `json:"workspaceId"` + ExploreMode bool `json:"exploreMode"` + ArtifactEnabled bool `json:"artifactEnabled"` } func (q *Queries) CreateOrUpdateChatSessionByUUID(ctx context.Context, arg CreateOrUpdateChatSessionByUUIDParams) (ChatSession, error) { @@ -535,28 +535,28 @@ ORDER BY ` type GetSessionsGroupedByWorkspaceRow struct { - ID int32 `json:"id"` - UserID int32 `json:"userId"` - Uuid string `json:"uuid"` - Topic string `json:"topic"` - CreatedAt time.Time `json:"createdAt"` - UpdatedAt time.Time `json:"updatedAt"` - Active bool `json:"active"` - Model string `json:"model"` - MaxLength int32 `json:"maxLength"` - Temperature float64 `json:"temperature"` - TopP float64 `json:"topP"` - MaxTokens int32 `json:"maxTokens"` - N int32 `json:"n"` - SummarizeMode bool `json:"summarizeMode"` - WorkspaceID sql.NullInt32 `json:"workspaceId"` - ArtifactEnabled bool `json:"artifactEnabled"` - Debug bool `json:"debug"` - ExploreMode bool `json:"exploreMode"` - WorkspaceUuid sql.NullString `json:"workspaceUuid"` - WorkspaceName sql.NullString `json:"workspaceName"` - WorkspaceColor sql.NullString `json:"workspaceColor"` - WorkspaceIcon sql.NullString `json:"workspaceIcon"` + ID int32 `json:"id"` + UserID int32 `json:"userId"` + Uuid string `json:"uuid"` + Topic string `json:"topic"` + CreatedAt time.Time `json:"createdAt"` + UpdatedAt time.Time `json:"updatedAt"` + Active bool `json:"active"` + Model string `json:"model"` + MaxLength int32 `json:"maxLength"` + Temperature float64 `json:"temperature"` + TopP float64 `json:"topP"` + MaxTokens int32 `json:"maxTokens"` + N int32 `json:"n"` + SummarizeMode bool `json:"summarizeMode"` + WorkspaceID sql.NullInt32 `json:"workspaceId"` + ArtifactEnabled bool `json:"artifactEnabled"` + Debug bool `json:"debug"` + ExploreMode bool `json:"exploreMode"` + WorkspaceUuid sql.NullString `json:"workspaceUuid"` + WorkspaceName sql.NullString `json:"workspaceName"` + WorkspaceColor sql.NullString `json:"workspaceColor"` + WorkspaceIcon sql.NullString `json:"workspaceIcon"` } func (q *Queries) GetSessionsGroupedByWorkspace(ctx context.Context, userID int32) ([]GetSessionsGroupedByWorkspaceRow, error) { diff --git a/api/sqlc_queries/models.go b/api/sqlc_queries/models.go index c64219fa..588a2ff5 100644 --- a/api/sqlc_queries/models.go +++ b/api/sqlc_queries/models.go @@ -132,24 +132,24 @@ type ChatPrompt struct { } type ChatSession struct { - ID int32 `json:"id"` - UserID int32 `json:"userId"` - Uuid string `json:"uuid"` - Topic string `json:"topic"` - CreatedAt time.Time `json:"createdAt"` - UpdatedAt time.Time `json:"updatedAt"` - Active bool `json:"active"` - Model string `json:"model"` - MaxLength int32 `json:"maxLength"` - Temperature float64 `json:"temperature"` - TopP float64 `json:"topP"` - MaxTokens int32 `json:"maxTokens"` - N int32 `json:"n"` - SummarizeMode bool `json:"summarizeMode"` - WorkspaceID sql.NullInt32 `json:"workspaceId"` - ArtifactEnabled bool `json:"artifactEnabled"` - Debug bool `json:"debug"` - ExploreMode bool `json:"exploreMode"` + ID int32 `json:"id"` + UserID int32 `json:"userId"` + Uuid string `json:"uuid"` + Topic string `json:"topic"` + CreatedAt time.Time `json:"createdAt"` + UpdatedAt time.Time `json:"updatedAt"` + Active bool `json:"active"` + Model string `json:"model"` + MaxLength int32 `json:"maxLength"` + Temperature float64 `json:"temperature"` + TopP float64 `json:"topP"` + MaxTokens int32 `json:"maxTokens"` + N int32 `json:"n"` + SummarizeMode bool `json:"summarizeMode"` + WorkspaceID sql.NullInt32 `json:"workspaceId"` + ArtifactEnabled bool `json:"artifactEnabled"` + Debug bool `json:"debug"` + ExploreMode bool `json:"exploreMode"` } type ChatSnapshot struct { diff --git a/api/svc/chat_auth_user_service.go b/api/svc/chat_auth_user_service.go index 37fbc071..d414b064 100644 --- a/api/svc/chat_auth_user_service.go +++ b/api/svc/chat_auth_user_service.go @@ -21,11 +21,11 @@ type AuthUserService struct { } // NewAuthUserService creates a new AuthUserService. -func NewAuthUserService(q *sqlc_queries.Queries) *AuthUserService { +func NewAuthUserService(q *sqlc_queries.Queries, jwtSecret string, defaultLimit int32) *AuthUserService { return &AuthUserService{ q: q, - jwtSecret: Cfg.JWTSecret, - defaultLimit: Cfg.DefaultLimit, + jwtSecret: jwtSecret, + defaultLimit: defaultLimit, } } diff --git a/api/svc/chat_main_service.go b/api/svc/chat_main_service.go index d0b9793e..3881e62f 100644 --- a/api/svc/chat_main_service.go +++ b/api/svc/chat_main_service.go @@ -5,7 +5,7 @@ import ( "context" "encoding/json" "fmt" - "log" + "log/slog" "net/http" "os" "strings" @@ -15,23 +15,25 @@ import ( "github.com/rotisserie/eris" "github.com/samber/lo" openai "github.com/sashabaranov/go-openai" + "github.com/swuecho/chat_backend/dto" "github.com/swuecho/chat_backend/llm/gemini" models "github.com/swuecho/chat_backend/models" "github.com/swuecho/chat_backend/provider" "github.com/swuecho/chat_backend/sqlc_queries" - "github.com/swuecho/chat_backend/dto" ) type ChatService struct { - q *sqlc_queries.Queries + q *sqlc_queries.Queries + openAIKey string + openAIProxy string } //go:embed artifact_instruction.txt var artifactInstructionText string -// NewChatService creates a new ChatService with database queries. -func NewChatService(q *sqlc_queries.Queries) *ChatService { - return &ChatService{q: q} +// NewChatService creates a new ChatService with database queries and OpenAI configuration. +func NewChatService(q *sqlc_queries.Queries, openAIKey, openAIProxy string) *ChatService { + return &ChatService{q: q, openAIKey: openAIKey, openAIProxy: openAIProxy} } // Q returns the underlying queries. @@ -125,7 +127,7 @@ func (s *ChatService) GetAskMessages(chatSession sqlc_queries.ChatSession, chatU if chatSession.ArtifactEnabled { artifactInstruction, err := LoadArtifactInstruction() if err != nil { - log.Printf("Warning: Failed to load artifact instruction: %v", err) + slog.Warn("Failed to load artifact instruction", "error", err) artifactInstruction = "" // Use empty string if file can't be loaded } @@ -140,7 +142,7 @@ func (s *ChatService) GetAskMessages(chatSession sqlc_queries.ChatSession, chatU func (s *ChatService) CreateChatPromptSimple(ctx context.Context, chatSessionUuid string, newQuestion string, userID int32) (sqlc_queries.ChatPrompt, error) { tokenCount, err := provider.GetTokenCount(newQuestion) if err != nil { - log.Printf("Warning: Failed to get token count for prompt: %v", err) + slog.Warn("Failed to get token count for prompt", "error", err) tokenCount = len(newQuestion) / dto.TokenEstimateRatio // Fallback estimate } chatPrompt, err := s.q.CreateChatPrompt(ctx, @@ -173,23 +175,23 @@ func (s *ChatService) CreateChatPromptSimple(ctx context.Context, chatSessionUui func (s *ChatService) CreateChatMessageSimple(ctx context.Context, sessionUuid, uuid, role, content, reasoningContent, model string, userId int32, baseURL string, is_summarize_mode bool) (sqlc_queries.ChatMessage, error) { numTokens, err := provider.GetTokenCount(content) if err != nil { - log.Printf("Warning: Failed to get token count: %v", err) + slog.Warn("Failed to get token count", "error", err) numTokens = len(content) / dto.TokenEstimateRatio // Fallback estimate } summary := "" if is_summarize_mode && numTokens > dto.SummarizeThreshold { - log.Println("summarizing") - summary = provider.SummarizeWithTimeout(baseURL, content) - log.Println("summarizing: " + summary) + slog.Info("summarizing") + summary = provider.SummarizeWithTimeout(s.openAIKey, baseURL, content) + slog.Info("summarizing: " + summary) } // Extract artifacts from content artifacts := extractArtifacts(content) artifactsJSON, err := json.Marshal(artifacts) if err != nil { - log.Printf("Warning: Failed to marshal artifacts: %v", err) + slog.Warn("Failed to marshal artifacts", "error", err) artifactsJSON = json.RawMessage([]byte("[]")) } @@ -220,22 +222,22 @@ func (s *ChatService) CreateChatMessageSimple(ctx context.Context, sessionUuid, func (s *ChatService) CreateChatMessageWithSuggestedQuestions(ctx context.Context, sessionUuid, uuid, role, content, reasoningContent, model string, userId int32, baseURL string, is_summarize_mode, exploreMode bool, messages []models.Message) (sqlc_queries.ChatMessage, error) { numTokens, err := provider.GetTokenCount(content) if err != nil { - log.Printf("Warning: Failed to get token count: %v", err) + slog.Warn("Failed to get token count", "error", err) numTokens = len(content) / dto.TokenEstimateRatio // Fallback estimate } summary := "" if is_summarize_mode && numTokens > dto.SummarizeThreshold { - log.Println("summarizing") - summary = provider.SummarizeWithTimeout(baseURL, content) - log.Println("summarizing: " + summary) + slog.Info("summarizing") + summary = provider.SummarizeWithTimeout(s.openAIKey, baseURL, content) + slog.Info("summarizing: " + summary) } // Extract artifacts from content artifacts := extractArtifacts(content) artifactsJSON, err := json.Marshal(artifacts) if err != nil { - log.Printf("Warning: Failed to marshal artifacts: %v", err) + slog.Warn("Failed to marshal artifacts", "error", err) artifactsJSON = json.RawMessage([]byte("[]")) } @@ -246,7 +248,7 @@ func (s *ChatService) CreateChatMessageWithSuggestedQuestions(ctx context.Contex if questionsJSON, err := json.Marshal(questions); err == nil { suggestedQuestions = questionsJSON } else { - log.Printf("Warning: Failed to marshal suggested questions: %v", err) + slog.Warn("Failed to marshal suggested questions", "error", err) } } @@ -322,7 +324,7 @@ func (s *ChatService) callLLMForSuggestions(prompt string) string { // Get all models and find preferred models for suggestions allModels, err := s.q.ListChatModels(ctx) if err != nil { - log.Printf("Warning: Failed to list models for suggestions: %v", err) + slog.Warn("Failed to list models for suggestions", "error", err) return "" } @@ -361,7 +363,7 @@ func (s *ChatService) callLLMForSuggestions(prompt string) string { } if selectedModel.ID == 0 { - log.Printf("Warning: No suitable models available for suggestions") + slog.Warn("No suitable models available for suggestions") return "" } @@ -375,7 +377,7 @@ func (s *ChatService) callLLMForSuggestions(prompt string) string { return s.callOpenAICompatibleForSuggestions(ctx, selectedModel, prompt) } - log.Printf("Warning: Unsupported model type for suggestions: %s", selectedModel.ApiType) + slog.Warn("Unsupported model type for suggestions", "apiType", selectedModel.ApiType) return "" } @@ -384,7 +386,7 @@ func (s *ChatService) callGeminiForSuggestions(ctx context.Context, model sqlc_q // Validate API key apiKey := os.Getenv("GEMINI_API_KEY") if apiKey == "" { - log.Printf("Warning: GEMINI_API_KEY environment variable not set") + slog.Warn("GEMINI_API_KEY environment variable not set") return "" } @@ -399,7 +401,7 @@ func (s *ChatService) callGeminiForSuggestions(ctx context.Context, model sqlc_q // Generate Gemini payload payloadBytes, err := gemini.GenGemminPayload(messages, nil) if err != nil { - log.Printf("Warning: Failed to generate Gemini payload for suggestions: %v", err) + slog.Warn("Failed to generate Gemini payload for suggestions", "error", err) return "" } @@ -407,7 +409,7 @@ func (s *ChatService) callGeminiForSuggestions(ctx context.Context, model sqlc_q url := gemini.BuildAPIURL(model.Name, false) req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(payloadBytes)) if err != nil { - log.Printf("Warning: Failed to create Gemini request for suggestions: %v", err) + slog.Warn("Failed to create Gemini request for suggestions", "error", err) return "" } req.Header.Set("Content-Type", "application/json") @@ -418,12 +420,12 @@ func (s *ChatService) callGeminiForSuggestions(ctx context.Context, model sqlc_q answer, err := gemini.HandleRegularResponse(http.Client{Timeout: 30 * time.Second}, req) if err != nil { - log.Printf("Warning: Failed to get Gemini response for suggestions: %v", err) + slog.Warn("Failed to get Gemini response for suggestions", "error", err) return "" } if answer == nil || answer.Answer == "" { - log.Printf("Warning: Empty response from Gemini for suggestions") + slog.Warn("Empty response from Gemini for suggestions") return "" } @@ -433,9 +435,9 @@ func (s *ChatService) callGeminiForSuggestions(ctx context.Context, model sqlc_q // callOpenAICompatibleForSuggestions makes an OpenAI-compatible API call for suggestions (including deepseek) func (s *ChatService) callOpenAICompatibleForSuggestions(ctx context.Context, model sqlc_queries.ChatModel, prompt string) string { // Generate OpenAI client configuration - config, err := provider.GenOpenAIConfig(model, provider.Config{OpenAIKey: Cfg.OpenAIKey, OpenAIProxy: Cfg.OpenAIProxy}) + config, err := provider.GenOpenAIConfig(model, provider.Config{OpenAIKey: s.openAIKey, OpenAIProxy: s.openAIProxy}) if err != nil { - log.Printf("Warning: Failed to generate OpenAI configuration for suggestions: %v", err) + slog.Warn("Failed to generate OpenAI configuration for suggestions", "error", err) return "" } @@ -460,12 +462,12 @@ func (s *ChatService) callOpenAICompatibleForSuggestions(ctx context.Context, mo resp, err := client.CreateChatCompletion(ctx, req) if err != nil { - log.Printf("Warning: Failed to generate suggested questions with %s: %v", model.Name, err) + slog.Warn("Failed to generate suggested questions", "model", model.Name, "error", err) return "" } if len(resp.Choices) == 0 { - log.Printf("Warning: No response choices returned for suggested questions from %s", model.Name) + slog.Warn("No response choices for suggested questions", "model", model.Name) return "" } @@ -479,7 +481,7 @@ func (s *ChatService) UpdateChatMessageContent(ctx context.Context, uuid, conten // num_tokens num_tokens, err := provider.GetTokenCount(content) if err != nil { - log.Printf("Warning: Failed to get token count for update: %v", err) + slog.Warn("Failed to get token count for update", "error", err) num_tokens = len(content) / dto.TokenEstimateRatio // Fallback estimate } @@ -506,17 +508,17 @@ func (s *ChatService) LogChat(chatSession sqlc_queries.ChatSession, msgs []model // log chat sessionRaw := chatSession.ToRawMessage() if sessionRaw == nil { - log.Println("failed to marshal chat session") + slog.Info("failed to marshal chat session") return } question, err := json.Marshal(msgs) if err != nil { - log.Printf("Warning: Failed to marshal chat messages: %v", err) + slog.Warn("Failed to marshal chat messages", "error", err) return // Skip logging if marshalling fails } answerRaw, err := json.Marshal(answerText) if err != nil { - log.Printf("Warning: Failed to marshal answer: %v", err) + slog.Warn("Failed to marshal answer", "error", err) return // Skip logging if marshalling fails } diff --git a/api/svc/chat_session_service.go b/api/svc/chat_session_service.go index 37d76353..06808264 100644 --- a/api/svc/chat_session_service.go +++ b/api/svc/chat_session_service.go @@ -8,9 +8,9 @@ import ( "github.com/rotisserie/eris" "github.com/samber/lo" + "github.com/swuecho/chat_backend/dto" "github.com/swuecho/chat_backend/provider" "github.com/swuecho/chat_backend/sqlc_queries" - "github.com/swuecho/chat_backend/dto" ) // ChatSessionService provides methods for interacting with chat sessions. diff --git a/api/svc/chat_snapshot_service.go b/api/svc/chat_snapshot_service.go index eed7d242..da375fe9 100644 --- a/api/svc/chat_snapshot_service.go +++ b/api/svc/chat_snapshot_service.go @@ -3,7 +3,7 @@ package svc import ( "context" "encoding/json" - "log" + "log/slog" "github.com/google/uuid" "github.com/samber/lo" @@ -81,7 +81,7 @@ func (s *ChatSnapshotService) CreateChatSnapshot(ctx context.Context, chatSessio Text: text, Conversation: simple_msgs_raw, }) if err != nil { - log.Println(err) + slog.Info("error", "error", err) return "", err } return one.Uuid, nil @@ -94,7 +94,7 @@ func GenTitle(q *sqlc_queries.Queries, ctx context.Context, chatSession sqlc_que if err == nil { genTitle, err := provider.GenerateChatTitle(ctx, model, text) if err != nil { - log.Println(err) + slog.Info("error", "error", err) } if genTitle != "" { title = genTitle @@ -131,7 +131,7 @@ func (s *ChatSnapshotService) CreateChatBot(ctx context.Context, chatSessionUuid Tags: json.RawMessage([]byte("{}")), Text: text, Conversation: simple_msgs_raw, }) if err != nil { - log.Println(err) + slog.Info("error", "error", err) return "", err } return one.Uuid, nil diff --git a/api/svc/chat_workspace_service.go b/api/svc/chat_workspace_service.go index 25330c8c..0eed66a5 100644 --- a/api/svc/chat_workspace_service.go +++ b/api/svc/chat_workspace_service.go @@ -4,7 +4,7 @@ import ( "context" "database/sql" "errors" - "log" + "log/slog" "github.com/google/uuid" "github.com/rotisserie/eris" @@ -116,7 +116,7 @@ func (s *ChatWorkspaceService) SetWorkspaceAsDefaultForUser(ctx context.Context, // --- Permission --- func (s *ChatWorkspaceService) HasWorkspacePermission(ctx context.Context, uuid string, userID int32) (bool, error) { - log.Printf("Checking permission for workspace=%s, user=%d", uuid, userID) + slog.Info("Checking permission for workspace=%s, user=%d", uuid, userID) result, err := s.q.HasWorkspacePermission(ctx, sqlc_queries.HasWorkspacePermissionParams{ Uuid: uuid, UserID: userID, }) @@ -210,7 +210,7 @@ func (s *ChatWorkspaceService) MigrateLegacyActiveSessions(ctx context.Context, ChatSessionUuid: session.ChatSessionUuid, }) if err != nil { - log.Printf("Warning: failed to migrate active session %s: %v", session.ChatSessionUuid, err) + slog.Warn("failed to migrate active session %s: %v", session.ChatSessionUuid, err) continue } // Delete old global active session diff --git a/api/svc/file_upload_service.go b/api/svc/file_upload_service.go index 666bfda6..c107fcf0 100644 --- a/api/svc/file_upload_service.go +++ b/api/svc/file_upload_service.go @@ -2,7 +2,7 @@ package svc import ( "context" - "log" + "log/slog" "github.com/swuecho/chat_backend/dto" "github.com/swuecho/chat_backend/sqlc_queries" @@ -37,15 +37,14 @@ func (s *ChatFileService) CreateChatUpload(ctx context.Context, params sqlc_quer return sqlc_queries.ChatFile{}, dto.ErrValidationInvalidInput("empty file data") } - log.Printf("Creating chat file upload for session %s, user %d", - params.ChatSessionUuid, params.UserID) + slog.Info("Creating chat file upload", "session", params.ChatSessionUuid, "userID", params.UserID) upload, err := s.q.CreateChatFile(ctx, params) if err != nil { return sqlc_queries.ChatFile{}, dto.WrapError(err, "failed to create chat file") } - log.Printf("Created chat file upload ID %d", upload.ID) + slog.Info("Created chat file upload", "id", upload.ID) return upload, nil } @@ -55,7 +54,7 @@ func (s *ChatFileService) GetChatFile(ctx context.Context, id int32) (sqlc_queri return sqlc_queries.GetChatFileByIDRow{}, dto.ErrValidationInvalidInput("invalid file ID") } - log.Printf("Retrieving chat file ID %d", id) + slog.Info("Retrieving chat file", "id", id) file, err := s.q.GetChatFileByID(ctx, id) if err != nil { @@ -71,7 +70,7 @@ func (s *ChatFileService) DeleteChatFile(ctx context.Context, id int32) error { return dto.ErrValidationInvalidInput("invalid file ID") } - log.Printf("Deleting chat file ID %d", id) + slog.Info("Deleting chat file", "id", id) _, err := s.q.DeleteChatFile(ctx, id) if err != nil { @@ -90,7 +89,7 @@ func (s *ChatFileService) ListChatFilesBySession(ctx context.Context, sessionUUI return nil, dto.ErrValidationInvalidInput("invalid user ID") } - log.Printf("Listing chat files for session %s, user %d", sessionUUID, userID) + slog.Info("Listing chat files", "session", sessionUUID, "userID", userID) files, err := s.q.ListChatFilesBySessionUUID(ctx, sqlc_queries.ListChatFilesBySessionUUIDParams{ ChatSessionUuid: sessionUUID, diff --git a/api/svc/service_error_wrapping_test.go b/api/svc/service_error_wrapping_test.go index c020f6c0..8e686798 100644 --- a/api/svc/service_error_wrapping_test.go +++ b/api/svc/service_error_wrapping_test.go @@ -27,7 +27,7 @@ func TestServiceErrorWrappingNil(t *testing.T) { q := sqlc_queries.New(testDB) ctx := context.Background() - testUser, err := NewAuthUserService(q).CreateAuthUser(ctx, + testUser, err := NewAuthUserService(q, "test-secret", 100).CreateAuthUser(ctx, sqlc_queries.CreateAuthUserParams{ Email: "wrap-test@test.com", Username: "wraptest", @@ -85,7 +85,7 @@ func TestServiceErrorWrappingNil(t *testing.T) { }) t.Run("AuthUserService_success_returns_nil_error", func(t *testing.T) { - svc := NewAuthUserService(q) + svc := NewAuthUserService(q, "test-secret", 100) stats, total, err := svc.GetUserStats(ctx, dto.Pagination{Page: 1, Size: 10}, 100) if err != nil { t.Fatalf("GetUserStats failed: %v", err) diff --git a/api/svc/svc.go b/api/svc/svc.go index 75d3477c..bfb2c4d1 100644 --- a/api/svc/svc.go +++ b/api/svc/svc.go @@ -1,10 +1,2 @@ // Package svc provides business logic services for the chat application. package svc - -// Cfg holds configuration set by main during initialization. -var Cfg struct { - OpenAIKey string - OpenAIProxy string - JWTSecret string - DefaultLimit int32 -} diff --git a/api/testutil/testdb.go b/api/testutil/testdb.go index 4a166da9..881ce419 100644 --- a/api/testutil/testdb.go +++ b/api/testutil/testdb.go @@ -5,6 +5,7 @@ import ( "database/sql" "fmt" "io" + "log" "os" "testing" "time" @@ -12,7 +13,6 @@ import ( _ "github.com/lib/pq" "github.com/ory/dockertest/v3" "github.com/ory/dockertest/v3/docker" - log "github.com/sirupsen/logrus" ) // NewTestDB spins up a PostgreSQL Docker container, runs schema migrations, diff --git a/api/util.go b/api/util.go index 17433cbc..a3d36039 100644 --- a/api/util.go +++ b/api/util.go @@ -1,36 +1,11 @@ package main import ( - "context" - "encoding/json" - "errors" - "net/http" - "os" - "strconv" "strings" - - "github.com/google/uuid" - "github.com/pkoukk/tiktoken-go" - "github.com/swuecho/chat_backend/dto" - "github.com/swuecho/chat_backend/middleware" ) -func NewUUID() string { - uuidv7, err := uuid.NewV7() - if err != nil { - return uuid.NewString() - } - return uuidv7.String() -} - -func getTokenCount(content string) (int, error) { - tke, err := tiktoken.GetEncoding("cl100k_base") - if err != nil { - return 0, err - } - return len(tke.Encode(content, nil, nil)), nil -} - +// firstN returns the first n characters (runes) of a string. +// Used by tests. func firstN(s string, n int) string { i := 0 for j := range s { @@ -42,6 +17,8 @@ func firstN(s string, n int) string { return s } +// firstNWords returns the first n words of a string. +// Used by tests. func firstNWords(s string, n int) string { if s == "" { return "" @@ -52,76 +29,3 @@ func firstNWords(s string, n int) string { } return strings.Join(words[:n], " ") } - -func setSSEHeader(w http.ResponseWriter) { - w.Header().Set("Content-Type", "text/event-stream") - w.Header().Set("Cache-Control", "no-cache, no-transform") - w.Header().Set("Connection", "keep-alive") - w.Header().Set("Access-Control-Allow-Origin", "*") - w.Header().Set("X-Accel-Buffering", "no") - w.Header().Del("Content-Length") - w.Header().Del("Content-Encoding") -} - -func setupSSEStream(w http.ResponseWriter) (http.Flusher, error) { - setSSEHeader(w) - flusher, ok := w.(http.Flusher) - if !ok { - return nil, errors.New(dto.ErrorStreamUnsupported) - } - return flusher, nil -} - -func getPerWordStreamLimit() int { - perWordStreamLimitStr := os.Getenv("PER_WORD_STREAM_LIMIT") - if perWordStreamLimitStr == "" { - return 200 - } - perWordStreamLimit, err := strconv.Atoi(perWordStreamLimitStr) - if err != nil { - return 200 - } - return perWordStreamLimit -} - -func getPaginationParams(r *http.Request) (limit int32, offset int32) { - limitStr := r.URL.Query().Get("limit") - offsetStr := r.URL.Query().Get("offset") - limit = 100 - if limitStr != "" { - if l, err := strconv.ParseInt(limitStr, 10, 32); err == nil { - limit = int32(l) - } - } - offset = 0 - if offsetStr != "" { - if o, err := strconv.ParseInt(offsetStr, 10, 32); err == nil { - offset = int32(o) - } - } - return limit, offset -} - -func getLimitParam(r *http.Request, defaultLimit int32) int32 { - limitStr := r.URL.Query().Get("limit") - if limitStr == "" { - return defaultLimit - } - limit, err := strconv.ParseInt(limitStr, 10, 32) - if err != nil { - return defaultLimit - } - return int32(limit) -} - -func DecodeJSON(r *http.Request, target interface{}) error { - return json.NewDecoder(r.Body).Decode(target) -} - -// getContextWithUser creates a background context with the given user ID. -// Used by test helpers. -func getContextWithUser(userID int) context.Context { - return context.WithValue(context.Background(), middleware.UserContextKey, strconv.Itoa(userID)) -} - -