mirror of
https://github.com/coder/coder.git
synced 2025-07-12 00:14:10 +00:00
feat: add database tables and API routes for agentic chat feature (#17570)
Backend portion of experimental `AgenticChat` feature: - Adds database tables for chats and chat messages - Adds functionality to stream messages from LLM providers using `kylecarbs/aisdk-go` - Adds API routes with relevant functionality (list, create, update chats, insert chat message) - Adds experiment `codersdk.AgenticChat` --------- Co-authored-by: Kyle Carberry <kyle@carberry.com>
This commit is contained in:
Cian Johnston
GitHub
@ -1269,6 +1269,10 @@ func (q *querier) DeleteApplicationConnectAPIKeysByUserID(ctx context.Context, u
|
||||
return q.db.DeleteApplicationConnectAPIKeysByUserID(ctx, userID)
|
||||
}
|
||||
|
||||
func (q *querier) DeleteChat(ctx context.Context, id uuid.UUID) error {
|
||||
return deleteQ(q.log, q.auth, q.db.GetChatByID, q.db.DeleteChat)(ctx, id)
|
||||
}
|
||||
|
||||
func (q *querier) DeleteCoordinator(ctx context.Context, id uuid.UUID) error {
|
||||
if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceTailnetCoordinator); err != nil {
|
||||
return err
|
||||
@ -1686,6 +1690,22 @@ func (q *querier) GetAuthorizationUserRoles(ctx context.Context, userID uuid.UUI
|
||||
return q.db.GetAuthorizationUserRoles(ctx, userID)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatByID(ctx context.Context, id uuid.UUID) (database.Chat, error) {
|
||||
return fetch(q.log, q.auth, q.db.GetChatByID)(ctx, id)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatMessagesByChatID(ctx context.Context, chatID uuid.UUID) ([]database.ChatMessage, error) {
|
||||
c, err := q.GetChatByID(ctx, chatID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return q.db.GetChatMessagesByChatID(ctx, c.ID)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatsByOwnerID(ctx context.Context, ownerID uuid.UUID) ([]database.Chat, error) {
|
||||
return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.GetChatsByOwnerID)(ctx, ownerID)
|
||||
}
|
||||
|
||||
func (q *querier) GetCoordinatorResumeTokenSigningKey(ctx context.Context) (string, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil {
|
||||
return "", err
|
||||
@ -3315,6 +3335,21 @@ func (q *querier) InsertAuditLog(ctx context.Context, arg database.InsertAuditLo
|
||||
return insert(q.log, q.auth, rbac.ResourceAuditLog, q.db.InsertAuditLog)(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) InsertChat(ctx context.Context, arg database.InsertChatParams) (database.Chat, error) {
|
||||
return insert(q.log, q.auth, rbac.ResourceChat.WithOwner(arg.OwnerID.String()), q.db.InsertChat)(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) InsertChatMessages(ctx context.Context, arg database.InsertChatMessagesParams) ([]database.ChatMessage, error) {
|
||||
c, err := q.db.GetChatByID(ctx, arg.ChatID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, c); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return q.db.InsertChatMessages(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) InsertCryptoKey(ctx context.Context, arg database.InsertCryptoKeyParams) (database.CryptoKey, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionCreate, rbac.ResourceCryptoKey); err != nil {
|
||||
return database.CryptoKey{}, err
|
||||
@ -3963,6 +3998,13 @@ func (q *querier) UpdateAPIKeyByID(ctx context.Context, arg database.UpdateAPIKe
|
||||
return update(q.log, q.auth, fetch, q.db.UpdateAPIKeyByID)(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) UpdateChatByID(ctx context.Context, arg database.UpdateChatByIDParams) error {
|
||||
fetch := func(ctx context.Context, arg database.UpdateChatByIDParams) (database.Chat, error) {
|
||||
return q.db.GetChatByID(ctx, arg.ID)
|
||||
}
|
||||
return update(q.log, q.auth, fetch, q.db.UpdateChatByID)(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) UpdateCryptoKeyDeletesAt(ctx context.Context, arg database.UpdateCryptoKeyDeletesAtParams) (database.CryptoKey, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceCryptoKey); err != nil {
|
||||
return database.CryptoKey{}, err
|
||||
|
||||
@ -5307,3 +5307,77 @@ func (s *MethodTestSuite) TestResourcesProvisionerdserver() {
|
||||
}).Asserts(rbac.ResourceWorkspaceAgentDevcontainers, policy.ActionCreate)
|
||||
}))
|
||||
}
|
||||
|
||||
func (s *MethodTestSuite) TestChat() {
|
||||
createChat := func(t *testing.T, db database.Store) (database.User, database.Chat, database.ChatMessage) {
|
||||
t.Helper()
|
||||
|
||||
usr := dbgen.User(t, db, database.User{})
|
||||
chat := dbgen.Chat(s.T(), db, database.Chat{
|
||||
OwnerID: usr.ID,
|
||||
})
|
||||
msg := dbgen.ChatMessage(s.T(), db, database.ChatMessage{
|
||||
ChatID: chat.ID,
|
||||
})
|
||||
|
||||
return usr, chat, msg
|
||||
}
|
||||
|
||||
s.Run("DeleteChat", s.Subtest(func(db database.Store, check *expects) {
|
||||
_, c, _ := createChat(s.T(), db)
|
||||
check.Args(c.ID).Asserts(c, policy.ActionDelete)
|
||||
}))
|
||||
|
||||
s.Run("GetChatByID", s.Subtest(func(db database.Store, check *expects) {
|
||||
_, c, _ := createChat(s.T(), db)
|
||||
check.Args(c.ID).Asserts(c, policy.ActionRead).Returns(c)
|
||||
}))
|
||||
|
||||
s.Run("GetChatMessagesByChatID", s.Subtest(func(db database.Store, check *expects) {
|
||||
_, c, m := createChat(s.T(), db)
|
||||
check.Args(c.ID).Asserts(c, policy.ActionRead).Returns([]database.ChatMessage{m})
|
||||
}))
|
||||
|
||||
s.Run("GetChatsByOwnerID", s.Subtest(func(db database.Store, check *expects) {
|
||||
u1, u1c1, _ := createChat(s.T(), db)
|
||||
u1c2 := dbgen.Chat(s.T(), db, database.Chat{
|
||||
OwnerID: u1.ID,
|
||||
CreatedAt: u1c1.CreatedAt.Add(time.Hour),
|
||||
})
|
||||
_, _, _ = createChat(s.T(), db) // other user's chat
|
||||
check.Args(u1.ID).Asserts(u1c2, policy.ActionRead, u1c1, policy.ActionRead).Returns([]database.Chat{u1c2, u1c1})
|
||||
}))
|
||||
|
||||
s.Run("InsertChat", s.Subtest(func(db database.Store, check *expects) {
|
||||
usr := dbgen.User(s.T(), db, database.User{})
|
||||
check.Args(database.InsertChatParams{
|
||||
OwnerID: usr.ID,
|
||||
Title: "test chat",
|
||||
CreatedAt: dbtime.Now(),
|
||||
UpdatedAt: dbtime.Now(),
|
||||
}).Asserts(rbac.ResourceChat.WithOwner(usr.ID.String()), policy.ActionCreate)
|
||||
}))
|
||||
|
||||
s.Run("InsertChatMessages", s.Subtest(func(db database.Store, check *expects) {
|
||||
usr := dbgen.User(s.T(), db, database.User{})
|
||||
chat := dbgen.Chat(s.T(), db, database.Chat{
|
||||
OwnerID: usr.ID,
|
||||
})
|
||||
check.Args(database.InsertChatMessagesParams{
|
||||
ChatID: chat.ID,
|
||||
CreatedAt: dbtime.Now(),
|
||||
Model: "test-model",
|
||||
Provider: "test-provider",
|
||||
Content: []byte(`[]`),
|
||||
}).Asserts(chat, policy.ActionUpdate)
|
||||
}))
|
||||
|
||||
s.Run("UpdateChatByID", s.Subtest(func(db database.Store, check *expects) {
|
||||
_, c, _ := createChat(s.T(), db)
|
||||
check.Args(database.UpdateChatByIDParams{
|
||||
ID: c.ID,
|
||||
Title: "new title",
|
||||
UpdatedAt: dbtime.Now(),
|
||||
}).Asserts(c, policy.ActionUpdate)
|
||||
}))
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user