Skip to content

Commit

Permalink
Improve a few things
Browse files Browse the repository at this point in the history
  • Loading branch information
@nono
nono committed Sep 23, 2024
1 parent 036ca1c commit 3f0bedd
Show file tree
Hide file tree
Showing 4 changed files with 74 additions and 47 deletions.
17 changes: 9 additions & 8 deletions docs/ai.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,11 +42,11 @@ vector database to find relevant documents (technically, only some parts of
the documents called chunks). Those documents are added to the prompt, so
that the LLM can use them as a context when answering.

### POST /ai/chat/completions/:id
### POST /ai/chat/conversations/:id

This route can be used to ask AI for a chat completion. The id in the path
must be the identifier of a chat session. The client can generate a random
identifier for a new chat session.
must be the identifier of a chat conversation. The client can generate a random
identifier for a new chat conversation.

The stack will respond after pushing a job for this task, but without the
response. The client must use the real-time websocket and subscribe to
Expand All @@ -55,7 +55,7 @@ response. The client must use the real-time websocket and subscribe to
#### Request

```http
POST /ai/chat/completions/e21dce8058b9013d800a18c04daba326 HTTP/1.1
POST /ai/chat/conversations/e21dce8058b9013d800a18c04daba326 HTTP/1.1
Content-Type: application/json
```

Expand All @@ -75,12 +75,13 @@ Content-Type: application/vnd.api+json
```json
{
"data": {
"type": "io.cozy.ai.chat.completions"
"type": "io.cozy.ai.chat.conversations"
"id": "e21dce8058b9013d800a18c04daba326",
"rev": "1-23456",
"attributes": {
"messages": [
{
"id": "eb17c3205bf1013ddea018c04daba326",
"role": "user",
"content": "Why the sky is blue?"
}
Expand All @@ -97,16 +98,16 @@ client > {"method": "AUTH", "payload": "token"}
client > {"method": "SUBSCRIBE",
"payload": {"type": "io.cozy.ai.chat.events"}}
server > {"event": "CREATED",
"payload": {"id": "e21dce8058b9013d800a18c04daba326",
"payload": {"id": "eb17c3205bf1013ddea018c04daba326",
"type": "io.cozy.ai.chat.events",
"doc": {"object": "delta", "content": "The "}}}
server > {"event": "CREATED",
"payload": {"id": "e21dce8058b9013d800a18c04daba326",
"payload": {"id": "eb17c3205bf1013ddea018c04daba326",
"type": "io.cozy.ai.chat.events",
"doc": {"object": "delta", "content": "sky "}}}
[...]
server > {"event": "CREATED",
"payload": {"id": "e21dce8058b9013d800a18c04daba326",
"payload": {"id": "eb17c3205bf1013ddea018c04daba326",
"type": "io.cozy.ai.chat.events",
"doc": {"object": "done"}}}
```
92 changes: 59 additions & 33 deletions model/rag/chat.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,69 +9,91 @@ import (
"io"
"net/http"
"net/url"
"time"

"github.com/cozy/cozy-stack/model/instance"
"github.com/cozy/cozy-stack/model/job"
"github.com/cozy/cozy-stack/pkg/consts"
"github.com/cozy/cozy-stack/pkg/couchdb"
"github.com/cozy/cozy-stack/pkg/jsonapi"
"github.com/cozy/cozy-stack/pkg/logger"
"github.com/cozy/cozy-stack/pkg/metadata"
"github.com/cozy/cozy-stack/pkg/realtime"
"github.com/gofrs/uuid/v5"
"github.com/labstack/echo/v4"
)

type ChatPayload struct {
ChatCompletionID string
Query string `json:"q"`
ChatConversationID string
Query string `json:"q"`
}

type ChatCompletion struct {
DocID string `json:"_id"`
DocRev string `json:"_rev,omitempty"`
Messages []ChatMessage `json:"messages"`
type ChatConversation struct {
DocID string `json:"_id"`
DocRev string `json:"_rev,omitempty"`
Messages []ChatMessage `json:"messages"`
Metadata *metadata.CozyMetadata `json:"cozyMetadata"`
}

type ChatMessage struct {
Role string `json:"role"`
Content string `json:"content"`
ID string `json:"id"`
Role string `json:"role"`
Content string `json:"content"`
CreatedAt time.Time
}

const (
HumanRole = "human"
AIRole = "ai"
UserRole = "user"
AssistantRole = "assistant"
)

func (c *ChatCompletion) ID() string { return c.DocID }
func (c *ChatCompletion) Rev() string { return c.DocRev }
func (c *ChatCompletion) DocType() string { return consts.ChatCompletions }
func (c *ChatCompletion) SetID(id string) { c.DocID = id }
func (c *ChatCompletion) SetRev(rev string) { c.DocRev = rev }
func (c *ChatCompletion) Clone() couchdb.Doc {
// DocTypeVersion represents the doctype version. Each time this document
// structure is modified, update this value
const DocTypeVersion = "1"

func (c *ChatConversation) ID() string { return c.DocID }
func (c *ChatConversation) Rev() string { return c.DocRev }
func (c *ChatConversation) DocType() string { return consts.ChatConversations }
func (c *ChatConversation) SetID(id string) { c.DocID = id }
func (c *ChatConversation) SetRev(rev string) { c.DocRev = rev }
func (c *ChatConversation) Clone() couchdb.Doc {
cloned := *c
cloned.Messages = make([]ChatMessage, len(c.Messages))
copy(cloned.Messages, c.Messages)
return &cloned
}
func (c *ChatCompletion) Included() []jsonapi.Object { return nil }
func (c *ChatCompletion) Relationships() jsonapi.RelationshipMap { return nil }
func (c *ChatCompletion) Links() *jsonapi.LinksList { return nil }
func (c *ChatConversation) Included() []jsonapi.Object { return nil }
func (c *ChatConversation) Relationships() jsonapi.RelationshipMap { return nil }
func (c *ChatConversation) Links() *jsonapi.LinksList { return nil }

var _ jsonapi.Object = (*ChatCompletion)(nil)
var _ jsonapi.Object = (*ChatConversation)(nil)

type QueryMessage struct {
Task string `json:"task"`
DocID string `json:"doc_id"`
}

func Chat(inst *instance.Instance, payload ChatPayload) (*ChatCompletion, error) {
var chat ChatCompletion
err := couchdb.GetDoc(inst, consts.ChatCompletions, payload.ChatCompletionID, &chat)
func Chat(inst *instance.Instance, payload ChatPayload) (*ChatConversation, error) {
var chat ChatConversation
err := couchdb.GetDoc(inst, consts.ChatConversations, payload.ChatConversationID, &chat)
if couchdb.IsNotFoundError(err) {
chat.DocID = payload.ChatCompletionID
chat.DocID = payload.ChatConversationID
md := metadata.New()
md.DocTypeVersion = DocTypeVersion
md.UpdatedAt = md.CreatedAt
chat.Metadata = md
} else if err != nil {
return nil, err
} else {
chat.Metadata.UpdatedAt = time.Now().UTC()
}
uuidv7, _ := uuid.NewV7()
msg := ChatMessage{
ID: uuidv7.String(),
Role: UserRole,
Content: payload.Query,
CreatedAt: time.Now().UTC(),
}
msg := ChatMessage{Role: HumanRole, Content: payload.Query}
chat.Messages = append(chat.Messages, msg)
if chat.DocRev == "" {
err = couchdb.CreateNamedDocWithDB(inst, &chat)
Expand Down Expand Up @@ -99,15 +121,14 @@ func Chat(inst *instance.Instance, payload ChatPayload) (*ChatCompletion, error)
}

func Query(inst *instance.Instance, logger logger.Logger, query QueryMessage) error {
var chat ChatCompletion
err := couchdb.GetDoc(inst, consts.ChatCompletions, query.DocID, &chat)
var chat ChatConversation
err := couchdb.GetDoc(inst, consts.ChatConversations, query.DocID, &chat)
if err != nil {
return err
}
msg := chat.Messages[len(chat.Messages)-1]
payload := map[string]interface{}{
"q": msg.Content,
"stream": true,
"messages": chat.Messages,
"stream": true,
}

res, err := callRAGQuery(inst, payload)
Expand All @@ -119,6 +140,7 @@ func Query(inst *instance.Instance, logger logger.Logger, query QueryMessage) er
return fmt.Errorf("POST status code: %d", res.StatusCode)
}

msg := chat.Messages[len(chat.Messages)-1]
var completion string
err = foreachSSE(res.Body, func(event map[string]interface{}) {
switch event["object"] {
Expand All @@ -129,6 +151,7 @@ func Query(inst *instance.Instance, logger logger.Logger, query QueryMessage) er
Type: consts.ChatEvents,
M: event,
}
delta.SetID(msg.ID)
go realtime.GetHub().Publish(inst, realtime.EventCreate, &delta, nil)
default:
// We can ignore done events
Expand All @@ -138,9 +161,12 @@ func Query(inst *instance.Instance, logger logger.Logger, query QueryMessage) er
return err
}

uuidv7, _ := uuid.NewV7()
answer := ChatMessage{
Role: AIRole,
Content: completion,
ID: uuidv7.String(),
Role: AssistantRole,
Content: completion,
CreatedAt: time.Now().UTC(),
}
chat.Messages = append(chat.Messages, answer)
return couchdb.UpdateDoc(inst, &chat)
Expand All @@ -155,7 +181,7 @@ func callRAGQuery(inst *instance.Instance, payload map[string]interface{}) (*htt
if err != nil {
return nil, err
}
u.Path = fmt.Sprintf("/query/%s", inst.Domain)
u.Path = fmt.Sprintf("/chat/%s", inst.Domain)
body, err := json.Marshal(payload)
if err != nil {
return nil, err
Expand Down
6 changes: 3 additions & 3 deletions pkg/consts/doctype.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -142,8 +142,8 @@ const (
// NextCloudFiles doc type is used when listing files from a NextCloud via
// WebDAV.
NextCloudFiles = "io.cozy.remote.nextcloud.files"
// ChatCompletions doc type is used for a chat between the user and a chatbot.
ChatCompletions = "io.cozy.ai.chat.completions"
// ChatEvents doc type is used for RAG events about a chat session.
// ChatConversations doc type is used for a chat between the user and a chatbot.
ChatConversations = "io.cozy.ai.chat.conversations"
// ChatEvents doc type is used for RAG events about a chat conversation.
ChatEvents = "io.cozy.ai.chat.events"
)
6 changes: 3 additions & 3 deletions web/ai/ai.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,14 +13,14 @@ import (

// Chat is the route for asking a chat completion to AI.
func Chat(c echo.Context) error {
if err := middlewares.AllowWholeType(c, permission.POST, consts.ChatCompletions); err != nil {
if err := middlewares.AllowWholeType(c, permission.POST, consts.ChatConversations); err != nil {
return middlewares.ErrForbidden
}
var payload rag.ChatPayload
if err := c.Bind(&payload); err != nil {
return err
}
payload.ChatCompletionID = c.Param("id")
payload.ChatConversationID = c.Param("id")
inst := middlewares.GetInstance(c)
chat, err := rag.Chat(inst, payload)
if err != nil {
Expand All @@ -31,5 +31,5 @@ func Chat(c echo.Context) error {

// Routes sets the routing for the AI tasks.
func Routes(router *echo.Group) {
router.POST("/chat/completions/:id", Chat)
router.POST("/chat/conversations/:id", Chat)
}

0 comments on commit 3f0bedd

Please sign in to comment.