llm chat support create assistant and thread

Makefile

@@ -57,6 +57,10 @@ migrate-down:
 	@echo "Migrating down..."
 	@migrate -path database/migrations -database "$(DATABASE_URL)" down
 
+psql:
+	@echo "Connecting to database..."
+	@docker compose exec -it postgres psql -U ${DATABASE_USER} -d ${DATABASE_NAME}
+
 sqlc:
 	@echo "Generating sqlc..."
 	@sqlc generate

database/queries/assistants.sql

@@ -1,8 +1,9 @@
 -- CRUD for assistants
 
--- name: CreateAssistant :exec
+-- name: CreateAssistant :one
 INSERT INTO assistants (user_id, name, description, system_prompt, model, metadata)
-VALUES ($1, $2, $3, $4, $5, $6);
+VALUES ($1, $2, $3, $4, $5, $6)
+RETURNING *;
 
 -- name: GetAssistant :one
 SELECT * FROM assistants WHERE uuid = $1;
@@ -18,9 +19,10 @@ DELETE FROM assistants WHERE uuid = $1;
 SELECT * FROM assistants WHERE user_id = $1 ORDER BY created_at DESC;
 
 -- CRUD for assistant_threads
--- name: CreateAssistantThread :exec
+-- name: CreateAssistantThread :one
 INSERT INTO assistant_threads (user_id, assistant_id, name, description, model, is_long_term_memory, metadata)
-VALUES ($1, $2, $3, $4, $5, $6, $7);
+VALUES ($1, $2, $3, $4, $5, $6, $7)
+RETURNING *;
 
 -- name: GetAssistantThread :one
 SELECT * FROM assistant_threads WHERE uuid = $1;
@@ -39,9 +41,10 @@ SELECT * FROM assistant_threads WHERE user_id = $1 ORDER BY created_at DESC;
 SELECT * FROM assistant_threads WHERE assistant_id = $1 ORDER BY created_at DESC;
 
 -- CRUD for assistant_thread_messages
--- name: CreateThreadMessage :exec
+-- name: CreateThreadMessage :one
 INSERT INTO assistant_messages (user_id, thread_id, model, token, role, text, attachments, metadata)
-VALUES ($1, $2, $3, $4, $5, $6, $7, $8);
+VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
+RETURNING *;
 
 -- name: GetThreadMessage :one
 SELECT * FROM assistant_messages WHERE uuid = $1;
@@ -56,9 +59,10 @@ DELETE FROM assistant_messages WHERE uuid = $1;
 SELECT * FROM assistant_messages WHERE thread_id = $1 ORDER BY created_at DESC;
 
 -- CRUD for assistant_attachments
--- name: CreateAttachment :exec
+-- name: CreateAttachment :one
 INSERT INTO assistant_attachments (user_id, entity, entity_id, file_type, file_url, size, metadata)
-VALUES ($1, $2, $3, $4, $5, $6, $7);
+VALUES ($1, $2, $3, $4, $5, $6, $7)
+RETURNING *;
 
 -- name: GetAttachment :one
 SELECT * FROM assistant_attachments WHERE uuid = $1;
@@ -76,9 +80,10 @@ SELECT * FROM assistant_attachments WHERE entity = $1 AND entity_id = $2 ORDER B
 SELECT * FROM assistant_attachments WHERE user_id = $1 ORDER BY created_at DESC;
 
 -- CRUD for assistant_message_embedddings
--- name: CreateAssistantEmbedding :exec
+-- name: CreateAssistantEmbedding :one
 INSERT INTO assistant_embedddings (user_id, message_id, attachment_id, embeddings)
-VALUES ($1, $2, $3, $4);
+VALUES ($1, $2, $3, $4)
+RETURNING *;
 
 -- name: DeleteAssistantEmbeddings :exec
 DELETE FROM assistant_embedddings WHERE id = $1;

database/queries/users.sql

@@ -1,13 +1,15 @@
 
--- name: InserUser :exec
+-- name: InserUser :one
 INSERT INTO users (username, telegram, activate_assistant_id, activate_thread_id)
-VALUES ($1, $2, $3, $4);
+VALUES ($1, $2, $3, $4)
+RETURNING *;
 
 -- name: GetTelegramUser :one
 SELECT * FROM users WHERE telegram = $1;
 
--- name: UpdateTelegramUser :exec
-UPDATE users SET activate_assistant_id = $1, activate_thread_id = $2 WHERE telegram = $3;
+-- name: UpdateTelegramUser :one
+UPDATE users SET activate_assistant_id = $1, activate_thread_id = $2 WHERE telegram = $3
+RETURNING *;
 
 -- name: DeleteTelegramUser :exec
 DELETE FROM users WHERE telegram = $1;

internal/core/assistants/assistant.go

@@ -0,0 +1,66 @@
+package assistants
+
+import (
+	"vibrain/internal/pkg/llms"
+
+	"github.com/google/uuid"
+)
+
+type AssistantMetaData struct{}
+
+type Assistant struct {
+	Id           uuid.UUID         `json:"uuid"`
+	UserId       uuid.UUID         `json:"user_id"`
+	Name         string            `json:"name"`
+	Description  string            `json:"description"`
+	SystemPrompt string            `json:"system_prompt"`
+	Model        string            `json:"model"`
+	MetaData     AssistantMetaData `json:"metadata"`
+}
+
+type AssistantOption func(*Assistant)
+
+func NewAssistant(userId uuid.UUID, opts ...AssistantOption) *Assistant {
+	a := &Assistant{
+		UserId:       userId,
+		Model:        llms.OpenAIGPT4oMini,
+		Name:         "Assistant" + uuid.New().String(),
+		Description:  "I am an AI assistant. I can help you with a variety of tasks.",
+		SystemPrompt: "You are a helpful AI assistant.",
+	}
+
+	for _, opt := range opts {
+		opt(a)
+	}
+	return a
+}
+
+func WithAssistantName(name string) AssistantOption {
+	return func(a *Assistant) {
+		a.Name = name
+	}
+}
+
+func WithAssistantDescription(description string) AssistantOption {
+	return func(a *Assistant) {
+		a.Description = description
+	}
+}
+
+func WithAssistantSystemPrompt(systemPrompt string) AssistantOption {
+	return func(a *Assistant) {
+		a.SystemPrompt = systemPrompt
+	}
+}
+
+func WithAssistantModel(model string) AssistantOption {
+	return func(a *Assistant) {
+		a.Model = model
+	}
+}
+
+func WithAssistantMetaData(metaData AssistantMetaData) AssistantOption {
+	return func(a *Assistant) {
+		a.MetaData = metaData
+	}
+}

internal/core/assistants/repository.go

@@ -0,0 +1,92 @@
+package assistants
+
+import (
+	"context"
+	"fmt"
+	"vibrain/internal/pkg/db"
+
+	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5/pgtype"
+)
+
+type Repository interface {
+	CreateAssistant(ctx context.Context, assistant *Assistant) error
+	GetAssistant(ctx context.Context, id uuid.UUID) (*Assistant, error)
+
+	CreateThread(ctx context.Context, thread *Thread) error
+	GetThread(ctx context.Context, id uuid.UUID) (*Thread, error)
+}
+
+type repository struct {
+	db *db.Queries
+}
+
+func NewRepository(pool *db.Pool) Repository {
+	return &repository{db: db.New(pool)}
+}
+
+func (r *repository) CreateAssistant(ctx context.Context, assistant *Assistant) error {
+	ast, err := r.db.CreateAssistant(ctx, db.CreateAssistantParams{
+		UserID:       pgtype.UUID{Bytes: assistant.UserId, Valid: true},
+		Name:         assistant.Name,
+		Description:  pgtype.Text{String: assistant.Description, Valid: assistant.Description != ""},
+		SystemPrompt: pgtype.Text{String: assistant.SystemPrompt, Valid: assistant.SystemPrompt != ""},
+		Model:        assistant.Model,
+	})
+	if err != nil {
+		return fmt.Errorf("failed to create assistant: %w", err)
+	}
+
+	assistant.Id = ast.Uuid
+	return nil
+}
+
+func (r *repository) GetAssistant(ctx context.Context, id uuid.UUID) (*Assistant, error) {
+	ast, err := r.db.GetAssistant(ctx, id)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get assistant: %w", err)
+	}
+
+	assistant := &Assistant{
+		Id:           ast.Uuid,
+		UserId:       ast.UserID.Bytes,
+		Name:         ast.Name,
+		Description:  ast.Description.String,
+		SystemPrompt: ast.SystemPrompt.String,
+		Model:        ast.Model,
+	}
+	return assistant, nil
+}
+
+func (r *repository) CreateThread(ctx context.Context, thread *Thread) error {
+	th, err := r.db.CreateAssistantThread(ctx, db.CreateAssistantThreadParams{
+		UserID:      pgtype.UUID{Bytes: thread.UserId, Valid: true},
+		AssistantID: pgtype.UUID{Bytes: thread.AssistantId, Valid: true},
+		Name:        thread.Name,
+		Description: pgtype.Text{String: thread.Description, Valid: thread.Description != ""},
+		Model:       thread.Model,
+	})
+	if err != nil {
+		return fmt.Errorf("failed to create thread: %w", err)
+	}
+
+	thread.Id = th.Uuid
+
+	return nil
+}
+
+func (r *repository) GetThread(ctx context.Context, id uuid.UUID) (*Thread, error) {
+	th, err := r.db.GetAssistantThread(ctx, id)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get thread: %w", err)
+	}
+	thread := &Thread{
+		Id:          th.Uuid,
+		UserId:      th.UserID.Bytes,
+		AssistantId: th.AssistantID.Bytes,
+		Name:        th.Name,
+		Description: th.Description.String,
+		Model:       th.Model,
+	}
+	return thread, nil
+}

internal/core/assistants/service.go

@@ -0,0 +1,36 @@
+package assistants
+
+import (
+	"context"
+	"vibrain/internal/pkg/db"
+
+	"github.com/google/uuid"
+)
+
+type Service struct {
+	db Repository
+}
+
+func NewService(db *db.Pool) (*Service, error) {
+	s := &Service{
+		db: NewRepository(db),
+	}
+
+	return s, nil
+}
+
+func (s *Service) CreateAssistant(ctx context.Context, assistant *Assistant) error {
+	return s.db.CreateAssistant(ctx, assistant)
+}
+
+func (s *Service) GetAssistant(ctx context.Context, id string) (*Assistant, error) {
+	return s.db.GetAssistant(ctx, uuid.MustParse(id))
+}
+
+func (s *Service) CreateThread(ctx context.Context, thread *Thread) error {
+	return s.db.CreateThread(ctx, thread)
+}
+
+func (s *Service) GetThread(ctx context.Context, id string) (*Thread, error) {
+	return s.db.GetThread(ctx, uuid.MustParse(id))
+}

internal/core/assistants/thread.go

@@ -0,0 +1,59 @@
+package assistants
+
+import (
+	"github.com/google/uuid"
+)
+
+type ThreadMetaData struct{}
+
+type Thread struct {
+	Id          uuid.UUID      `json:"uuid"`
+	UserId      uuid.UUID      `json:"user_id"`
+	AssistantId uuid.UUID      `json:"assistant_id"`
+	Name        string         `json:"name"`
+	Description string         `json:"description"`
+	Model       string         `json:"model"`
+	MetaData    ThreadMetaData `json:"metadata"`
+}
+
+type ThreadOption func(*Thread)
+
+func NewThread(userId uuid.UUID, assistant Assistant, opts ...ThreadOption) *Thread {
+	t := &Thread{
+		UserId:      userId,
+		AssistantId: assistant.Id,
+		Model:       assistant.Model,
+		Name:        "Thread" + uuid.New().String(),
+		Description: "I am a conversation thread.",
+	}
+
+	for _, opt := range opts {
+		opt(t)
+	}
+	return t
+}
+
+func WithThreadName(name string) ThreadOption {
+	return func(t *Thread) {
+		t.Name = name
+	}
+}
+
+func WithThreadDescription(description string) ThreadOption {
+	return func(t *Thread) {
+		t.Description = description
+	}
+}
+
+func WithThreadModel(model string) ThreadOption {
+	return func(t *Thread) {
+		t.Model = model
+	}
+}
+
+func WithThreadMetaData(metaData ThreadMetaData) ThreadOption {
+	return func(t *Thread) {
+		t.MetaData = metaData
+	}
+}
+