Skip to content

feat: add max tokens param for ai generation #44

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
elliot14A wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: add max tokens param for ai generation #44

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions server/application/repositories/ai.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,12 +7,12 @@ import (
)

type AiRepository interface {
GenerateSql(nl string) (string, error)
GenerateColumnDescriptions(ctx context.Context, rows string, summary string) (map[string]string, error)
GenerateDatasetDescription(ctx context.Context, datasetName string, columnNames []string, columnDescriptions map[string]string, rows string, summary string) (string, error)
GenerateSql(nl string, maxTokens *int) (string, error)
GenerateColumnDescriptions(ctx context.Context, rows string, summary string, maxTokens *int) (map[string]string, error)
GenerateDatasetDescription(ctx context.Context, datasetName string, columnNames []string, columnDescriptions map[string]string, rows string, summary string, maxTokens *int) (string, error)
}

type AiChatRepository interface {
GenerateChatResponse(ctx context.Context, userMessage string, prevMessage []*models.D_ChatMessage) (*models.D_AiChatResponse, error)
GenerateTitle(ctx context.Context, userMessage string) (*models.D_AiChatResponse, error)
GenerateChatResponse(ctx context.Context, userMessage string, prevMessage []*models.D_ChatMessage, maxTokens *int) (*models.D_AiChatResponse, error)
GenerateTitle(ctx context.Context, userMessage string, maxTokens *int) (*models.D_AiChatResponse, error)
}
Comment on lines 9 to 18
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot Dec 31, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Find implementations of AiRepository and AiChatRepository interface methods
# Check for GenerateSql, GenerateColumnDescriptions, GenerateDatasetDescription,
# GenerateChatResponse, and GenerateTitle implementations

echo "=== Searching for GenerateSql implementations ==="
rg "func.*GenerateSql" --type go -A 2

echo -e "\n=== Searching for GenerateColumnDescriptions implementations ==="
rg "func.*GenerateColumnDescriptions" --type go -A 2

echo -e "\n=== Searching for GenerateDatasetDescription implementations ==="
rg "func.*GenerateDatasetDescription" --type go -A 2

echo -e "\n=== Searching for GenerateChatResponse implementations ==="
rg "func.*GenerateChatResponse" --type go -A 2

echo -e "\n=== Searching for GenerateTitle implementations ==="
rg "func.*GenerateTitle" --type go -A 2

Repository: factly/gopie

Length of output: 3439

🏁 Script executed:

#!/bin/bash
# Search specifically for AiDriver implementations
echo "=== All GenerateChatResponse methods ==="
rg "GenerateChatResponse" --type go -B 1 -A 3

echo -e "\n=== All GenerateTitle methods ==="
rg "GenerateTitle" --type go -B 1 -A 3

Repository: factly/gopie

Length of output: 3669

🏁 Script executed:

cat -n server/application/services/ai.go

Repository: factly/gopie

Length of output: 1139

Fix critical signature mismatches in AiDriver implementations and add missing AiChatRepository methods.

The interface definitions were updated but the AiDriver wrapper implementation in server/application/services/ai.go has not been fully updated:

  1. Signature mismatch: GenerateColumnDescriptions and GenerateDatasetDescription are missing the ctx context.Context parameter required by the AiRepository interface. These methods pass context.Background() when calling the underlying implementation, but their own signatures don't accept a context parameter.

  2. Missing implementations: GenerateChatResponse and GenerateTitle methods from the AiChatRepository interface are not implemented in AiDriver at all. These are called from server/application/services/chat.go, so the AiDriver must implement them.

Update AiDriver to match both interface signatures exactly and implement all required methods from both AiRepository and AiChatRepository interfaces.

🤖 Prompt for AI Agents 
In server/application/repositories/ai.go around lines 9 to 18, the AiRepository
and AiChatRepository signatures were changed but the AiDriver wrapper in
server/application/services/ai.go was not updated: add the missing ctx
context.Context parameter to AiDriver.GenerateColumnDescriptions and
AiDriver.GenerateDatasetDescription so their signatures match the interface and
forward that ctx when calling the underlying implementation (remove hardcoded
context.Background()), and implement the two missing AiChatRepository methods
GenerateChatResponse(ctx context.Context, userMessage string, prevMessage
[]*models.D_ChatMessage, maxTokens *int) and GenerateTitle(ctx context.Context,
userMessage string, maxTokens *int) in AiDriver to delegate to the underlying ai
implementation, passing all parameters and returning its results/errors.

12 changes: 6 additions & 6 deletions server/application/services/ai.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,14 +14,14 @@ func NewAiDriver(ai repositories.AiRepository) *AiDriver {
return &AiDriver{ai}
}

func (d *AiDriver) GenerateSql(query string) (string, error) {
return d.ai.GenerateSql(query)
func (d *AiDriver) GenerateSql(query string, maxTokens *int) (string, error) {
return d.ai.GenerateSql(query, maxTokens)
}

func (d *AiDriver) GenerateColumnDescriptions(rows string, summary string) (map[string]string, error) {
return d.ai.GenerateColumnDescriptions(context.Background(), rows, summary)
func (d *AiDriver) GenerateColumnDescriptions(rows string, summary string, maxTokens *int) (map[string]string, error) {
return d.ai.GenerateColumnDescriptions(context.Background(), rows, summary, maxTokens)
}

func (d *AiDriver) GenerateDatasetDescription(datasetName string, columnNames []string, columnDescriptions map[string]string, rows string, summary string) (string, error) {
return d.ai.GenerateDatasetDescription(context.Background(), datasetName, columnNames, columnDescriptions, rows, summary)
func (d *AiDriver) GenerateDatasetDescription(datasetName string, columnNames []string, columnDescriptions map[string]string, rows string, summary string, maxTokens *int) (string, error) {
return d.ai.GenerateDatasetDescription(context.Background(), datasetName, columnNames, columnDescriptions, rows, summary, maxTokens)
}
10 changes: 5 additions & 5 deletions server/application/services/chat.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,11 +47,11 @@ func (service *ChatService) AddNewMessage(ctx context.Context, chatID string, me
return service.store.AddNewMessage(ctx, chatID, messages)
}

func (service *ChatService) D_ChatWithAi(params *models.D_ChatWithAiParams) (*models.D_ChatWithMessages, error) {
func (service *ChatService) D_ChatWithAi(params *models.D_ChatWithAiParams, maxTokens *int) (*models.D_ChatWithMessages, error) {
messages := params.Messages
prevMsgs := &models.PaginationView[*models.D_ChatMessage]{}

aiResponse, err := service.ai.GenerateChatResponse(context.Background(), params.Prompt, prevMsgs.Results)
aiResponse, err := service.ai.GenerateChatResponse(context.Background(), params.Prompt, prevMsgs.Results, maxTokens)
if err != nil {
return nil, fmt.Errorf("Error generating chat response from ai: %v", err)
}
Expand All @@ -64,7 +64,7 @@ func (service *ChatService) D_ChatWithAi(params *models.D_ChatWithAiParams) (*mo
// new chat
if params.ChatID == "" {

title, err := service.ai.GenerateTitle(context.Background(), messages[len(messages)-1].Content)
title, err := service.ai.GenerateTitle(context.Background(), messages[len(messages)-1].Content, maxTokens)
if err != nil {
return nil, fmt.Errorf("Error generating title from ai: %v", err)
}
Expand Down Expand Up @@ -94,7 +94,7 @@ func (service *ChatService) ChatWithAiAgent(ctx context.Context, params *models.
service.aiAgent.Chat(ctx, params)
}

func (service *ChatService) CreateChat(ctx context.Context, params *models.CreateChatParams) (*models.ChatWithMessages, error) {
func (service *ChatService) CreateChat(ctx context.Context, params *models.CreateChatParams, maxTokens *int) (*models.ChatWithMessages, error) {
var userMessage *models.ChatMessage
var filteredMessages []models.ChatMessage

Expand Down Expand Up @@ -127,7 +127,7 @@ func (service *ChatService) CreateChat(ctx context.Context, params *models.Creat
// Generate title only if we have a user message
if userMessage != nil && userMessage.Choices != nil && len(userMessage.Choices) > 0 &&
userMessage.Choices[0].Delta.Content != nil {
title, err := service.ai.GenerateTitle(ctx, *userMessage.Choices[0].Delta.Content)
title, err := service.ai.GenerateTitle(ctx, *userMessage.Choices[0].Delta.Content, maxTokens)
if err != nil {
fmt.Printf("Error generating title from AI: %v\n", err)
title = &models.D_AiChatResponse{
Expand Down
58 changes: 39 additions & 19 deletions server/infrastructure/openai/openai.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,8 @@ func NewOpenAIClient(cfg config.OpenAIConfig, logger *logger.Logger) *OpenAIClie
// GenerateResponseJSON is a generic function that generates a response from OpenAI
// and unmarshals it into the provided type T. It enforces JSON output format
// to ensure consistent structured responses.
func GenerateResponseJSON[T any](c *OpenAIClient, content string) (*T, error) {
// maxTokens is optional - pass nil to use default, or a pointer to an int to limit tokens
func GenerateResponseJSON[T any](c *OpenAIClient, content string, maxTokens *int) (*T, error) {
c.logger.Debug("generating JSON response from OpenAI", zap.String("model", c.model))
msgs := openai.ChatCompletionMessage{
Role: "user",
Expand All @@ -77,13 +78,19 @@ func GenerateResponseJSON[T any](c *OpenAIClient, content string) (*T, error) {

ctx := context.Background()

res, err := c.client.CreateChatCompletion(ctx, openai.ChatCompletionRequest{
req := openai.ChatCompletionRequest{
Model: c.model,
Messages: []openai.ChatCompletionMessage{msgs},
ResponseFormat: &openai.ChatCompletionResponseFormat{
Type: openai.ChatCompletionResponseFormatTypeJSONObject,
},
})
}

if maxTokens != nil {
req.MaxTokens = *maxTokens
}

res, err := c.client.CreateChatCompletion(ctx, req)
if err != nil {
c.logger.Error("failed to generate response from OpenAI", zap.Error(err))
return nil, err
Expand Down Expand Up @@ -111,7 +118,8 @@ func GenerateResponseJSON[T any](c *OpenAIClient, content string) (*T, error) {

// GenerateResponseString is a non-generic version for backwards compatibility
// when you need a plain string response without JSON parsing
func (c *OpenAIClient) GenerateResponseString(content string) (string, error) {
// maxTokens is optional - pass nil to use default, or a pointer to an int to limit tokens
func (c *OpenAIClient) GenerateResponseString(content string, maxTokens *int) (string, error) {
c.logger.Debug("generating string response from OpenAI", zap.String("model", c.model))
msgs := openai.ChatCompletionMessage{
Role: "user",
Expand All @@ -120,10 +128,16 @@ func (c *OpenAIClient) GenerateResponseString(content string) (string, error) {

ctx := context.Background()

res, err := c.client.CreateChatCompletion(ctx, openai.ChatCompletionRequest{
req := openai.ChatCompletionRequest{
Model: c.model,
Messages: []openai.ChatCompletionMessage{msgs},
})
}

if maxTokens != nil {
req.MaxTokens = *maxTokens
}

res, err := c.client.CreateChatCompletion(ctx, req)
if err != nil {
c.logger.Error("failed to generate response from OpenAI", zap.Error(err))
return "", err
Expand All @@ -141,7 +155,7 @@ func (c *OpenAIClient) GenerateResponseString(content string) (string, error) {
return responseContent, nil
}

func (c *OpenAIClient) GenerateChatResponseFunc(userMsg string, prevMsgs []*models.D_ChatMessage) (string, error) {
func (c *OpenAIClient) GenerateChatResponseFunc(userMsg string, prevMsgs []*models.D_ChatMessage, maxTokens *int) (string, error) {
c.logger.Debug("generating chat response from OpenAI",
zap.Int("previous_messages", len(prevMsgs)),
zap.String("model", c.model))
Expand All @@ -162,10 +176,16 @@ func (c *OpenAIClient) GenerateChatResponseFunc(userMsg string, prevMsgs []*mode
msgs = append(msgs, latestMessage)

ctx := context.Background()
res, err := c.client.CreateChatCompletion(ctx, openai.ChatCompletionRequest{
req := openai.ChatCompletionRequest{
Model: c.model,
Messages: msgs,
})
}

if maxTokens != nil {
req.MaxTokens = *maxTokens
}

res, err := c.client.CreateChatCompletion(ctx, req)
if err != nil {
c.logger.Error("failed to generate chat response from OpenAI", zap.Error(err))
return "", err
Expand All @@ -183,13 +203,13 @@ func (c *OpenAIClient) GenerateChatResponseFunc(userMsg string, prevMsgs []*mode
return responseContent, nil
}

func (c *OpenAIClient) GenerateSql(content string) (string, error) {
func (c *OpenAIClient) GenerateSql(content string, maxTokens *int) (string, error) {
c.logger.Debug("generating SQL from OpenAI")
return c.GenerateResponseString(content)
return c.GenerateResponseString(content, maxTokens)
}

func (c *OpenAIClient) GenerateChatResponse(ctx context.Context, userMessage string, prevMessages []*models.D_ChatMessage) (*models.D_AiChatResponse, error) {
resp, err := c.GenerateChatResponseFunc(userMessage, prevMessages)
func (c *OpenAIClient) GenerateChatResponse(ctx context.Context, userMessage string, prevMessages []*models.D_ChatMessage, maxTokens *int) (*models.D_AiChatResponse, error) {
resp, err := c.GenerateChatResponseFunc(userMessage, prevMessages, maxTokens)
if err != nil {
return nil, err
}
Expand All @@ -199,14 +219,14 @@ func (c *OpenAIClient) GenerateChatResponse(ctx context.Context, userMessage str
}, nil
}

func (c *OpenAIClient) GenerateTitle(ctx context.Context, content string) (*models.D_AiChatResponse, error) {
func (c *OpenAIClient) GenerateTitle(ctx context.Context, content string, maxTokens *int) (*models.D_AiChatResponse, error) {
c.logger.Debug("generating title from OpenAI")
systemPrompt := `
!! IMPORTANT: In the response only provide the title of the content. Do not provide any other information. !!
Generate a title for the following content:
` + content

resp, err := c.GenerateResponseString(systemPrompt)
resp, err := c.GenerateResponseString(systemPrompt, maxTokens)
if err != nil {
c.logger.Error("failed to generate title", zap.Error(err))
return nil, err
Expand All @@ -218,7 +238,7 @@ func (c *OpenAIClient) GenerateTitle(ctx context.Context, content string) (*mode
}, nil
}

func (c *OpenAIClient) GenerateColumnDescriptions(ctx context.Context, rows string, summary string) (map[string]string, error) {
func (c *OpenAIClient) GenerateColumnDescriptions(ctx context.Context, rows string, summary string, maxTokens *int) (map[string]string, error) {
c.logger.Debug("generating column descriptions from OpenAI")

const maxRetries = 3
Expand Down Expand Up @@ -264,7 +284,7 @@ Please generate a valid JSON response following the exact format specified above
zap.String("previous_error", lastError.Error()))
}

result, lastError = GenerateResponseJSON[map[string]string](c, systemPrompt)
result, lastError = GenerateResponseJSON[map[string]string](c, systemPrompt, maxTokens)
if lastError == nil && result != nil && len(*result) > 0 {
c.logger.Debug("successfully generated column descriptions",
zap.Int("columns_count", len(*result)),
Expand Down Expand Up @@ -302,7 +322,7 @@ Please generate a valid JSON response following the exact format specified above
return nil, fmt.Errorf("failed to generate column descriptions after %d attempts: %w", maxRetries, lastError)
}

func (c *OpenAIClient) GenerateDatasetDescription(ctx context.Context, datasetName string, columnNames []string, columnDescriptions map[string]string, rows string, summary string) (string, error) {
func (c *OpenAIClient) GenerateDatasetDescription(ctx context.Context, datasetName string, columnNames []string, columnDescriptions map[string]string, rows string, summary string, maxTokens *int) (string, error) {
c.logger.Debug("generating dataset description from OpenAI",
zap.String("dataset_name", datasetName),
zap.Int("column_count", len(columnNames)))
Expand Down Expand Up @@ -340,7 +360,7 @@ func (c *OpenAIClient) GenerateDatasetDescription(ctx context.Context, datasetNa
- Provide ONLY the description text, no additional formatting or explanations
`, datasetName, columnInfo, rows, summary)

resp, err := c.GenerateResponseString(systemPrompt)
resp, err := c.GenerateResponseString(systemPrompt, maxTokens)
if err != nil {
c.logger.Error("failed to generate dataset description", zap.Error(err))
return "", err
Expand Down
7 changes: 6 additions & 1 deletion server/interfaces/http/routes/api/ai/columns.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,8 @@ type genColumnsDescBody struct {
Summary any `json:"summary" validate:"required" example:"{\"column_name\": {\"type\": \"string\", \"count\": 1000}}"`
// Sample rows from the dataset to help AI understand the data context
Rows any `json:"rows" validate:"required" example:"[[\"value1\", \"value2\"], [\"value3\", \"value4\"]]"`
// Maximum number of tokens to generate (optional)
MaxTokens *int `json:"maxTokens,omitempty" example:"1000"`
}

// @Summary Generate AI-powered column descriptions
Expand All @@ -38,7 +40,7 @@ func (h *httpHandler) genColumnsDesc(c *fiber.Ctx) error {
rowsString := fmt.Sprintf("%v", body.Rows)
SummaryString := fmt.Sprintf("%v", body.Summary)

descriptions, err := h.aiSvc.GenerateColumnDescriptions(rowsString, SummaryString)
descriptions, err := h.aiSvc.GenerateColumnDescriptions(rowsString, SummaryString, body.MaxTokens)
if err != nil {
h.logger.Error("Error generating column descriptions", zap.Error(err))
return fiber.NewError(fiber.StatusInternalServerError, "Failed to generate column descriptions")
Expand All @@ -61,6 +63,8 @@ type genDatasetDescBody struct {
Rows any `json:"rows" validate:"required" example:"[[\"2024-01-01\", \"Widget A\", 10, 100.00]]"`
// Dataset summary statistics
Summary any `json:"summary" example:"{\"rowCount\": 1000, \"columnCount\": 4}"`
// Maximum number of tokens to generate (optional)
MaxTokens *int `json:"maxTokens,omitempty" example:"1000"`
}

// @Summary Generate AI-powered dataset description
Expand Down Expand Up @@ -89,6 +93,7 @@ func (h *httpHandler) genDatasetDesc(c *fiber.Ctx) error {
body.ColumnDescriptions,
fmt.Sprintf("%v", body.Rows),
fmt.Sprintf("%v", body.Summary),
body.MaxTokens,
)
if err != nil {
h.logger.Error("Error generating dataset description", zap.Error(err))
Expand Down
6 changes: 5 additions & 1 deletion server/interfaces/http/routes/api/chats/create.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -342,12 +342,16 @@ func (h *httpHandler) chatWithAgent(ctx *fiber.Ctx) error {
var chatWithMessages *models.ChatWithMessages
var err error
if chatIdHeader == "" {
var maxTokens *int
if body.MaxTokens > 0 {
maxTokens = &body.MaxTokens
}
chatWithMessages, err = h.chatSvc.CreateChat(context.Background(), &models.CreateChatParams{
ID: sessionID,
Messages: messages,
CreatedBy: userID,
OrganizationID: orgID,
})
}, maxTokens)
if err != nil {
h.logger.Critical("SSE: Error creating new chat", zap.Error(err), zap.String("session_id", sessionID))
errorEvent := pkg.ChatMessageFromError(err)
Expand Down
3 changes: 2 additions & 1 deletion server/interfaces/http/routes/api/chats/create_chat.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ type createChatRequestBody struct {
Title string `json:"title,omitempty"`
DatasetIDs []string `json:"dataset_ids,omitempty"`
ProjectIDs []string `json:"project_ids,omitempty"`
MaxTokens *int `json:"maxTokens,omitempty" example:"1000"`
}

// @Description Response body for creating a new chat session
Expand Down Expand Up @@ -79,7 +80,7 @@ func (h *httpHandler) createChat(ctx *fiber.Ctx) error {
}

// Use the CreateChat service which now handles empty messages
chatWithMessages, err := h.chatSvc.CreateChat(ctx.Context(), chatParams)
chatWithMessages, err := h.chatSvc.CreateChat(ctx.Context(), chatParams, body.MaxTokens)
if err != nil {
h.logger.Error("Error creating new chat", zap.Error(err), zap.String("chat_id", chatID))
return ctx.Status(fiber.StatusInternalServerError).JSON(fiber.Map{
Expand Down
4 changes: 3 additions & 1 deletion server/interfaces/http/routes/api/nl2sql.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@ type nl2SqlRequest struct {
Query string `json:"query" validate:"required,min=3" example:"show me total sales by region"`
// Name of the dataset/table to query
TableName string `json:"table" validate:"required" example:"sales_data"`
// Maximum number of tokens to generate (optional)
MaxTokens *int `json:"maxTokens,omitempty" example:"1000"`
}

// @Summary Convert natural language to SQL
Expand Down Expand Up @@ -120,7 +122,7 @@ func (h *httpHandler) nl2sql(ctx *fiber.Ctx) error {
- String only generate SQL for read queries nothing else.
`, body.Query, body.TableName, schemaJson, rowsCsv)

sql, err := h.aiSvc.GenerateSql(content)
sql, err := h.aiSvc.GenerateSql(content, body.MaxTokens)
if err != nil {
h.logger.Error("Error generating SQL", zap.Error(err))
return ctx.Status(fiber.StatusInternalServerError).JSON(fiber.Map{
Expand Down
5 changes: 4 additions & 1 deletion server/interfaces/http/routes/source/database/create.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,8 @@ type createRequestBody struct {
// Alias of the dataset
Alias string `json:"alias" validate:"required,min=3" example:"users_data"`
CustomPrompt string `json:"custom_prompt"`
// Maximum number of tokens to generate (optional)
MaxTokens *int `json:"maxTokens,omitempty" example:"1000"`
}

// @Summary Create dataset from Postgres
Expand Down Expand Up @@ -210,7 +212,7 @@ func (h *httpHandler) create(ctx *fiber.Ctx) error {
rowsBytes, _ := json.Marshal(rows)
rowsString := string(rowsBytes)

descriptions, err := h.aiSvc.GenerateColumnDescriptions(rowsString, summaryString)
descriptions, err := h.aiSvc.GenerateColumnDescriptions(rowsString, summaryString, body.MaxTokens)
if err != nil {
h.logger.Error("Error generating column descriptions", zap.Error(err))
h.cleanupResources(cleanup)
Expand All @@ -236,6 +238,7 @@ func (h *httpHandler) create(ctx *fiber.Ctx) error {
descriptions,
rowsString,
summaryString,
body.MaxTokens,
)
if err != nil {
h.logger.Error("Error generating dataset description", zap.Error(err))
Expand Down