Leaf/rag
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

改进 block

Browse Source
This commit is contained in:
ivamp 2024-07-21 17:24:56 +08:00
parent d6cffe911c
commit 16957ef1b3
13 changed files with 441 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
api/rag/document.proto
View File

@ -26,9 +26,38 @@ message ListDocumentResponse {
message CreateDocumentBlockRequest {
int64 document_id = 1;
int64 order = 2;
string content = 3;
}
message CreateDocumentBlockResponse {
DocumentBlock document_block = 1;
}
message DeleteDocumentBlockRequest {
int64 id = 1;
}
message DeleteDocumentBlockResponse {
}
message ListDocumentBlockRequest {
int64 document_id = 1;
}
message ListDocumentBlockResponse {
repeated DocumentBlock document_blocks = 1;
}
message UpdateDocumentBlockRequest {
int64 id = 1;
string content = 2;
}
message UpdateDocumentBlockResponse {
DocumentBlock document_block = 1;
}
message Document {
int64 id = 1;
string name = 2;

24
api/rag/rag.proto
View File

@ -37,5 +37,29 @@ service RAGService {
};
}
rpc ListDocumentBlock(ListDocumentBlockRequest) returns (ListDocumentBlockResponse) {
option (google.api.http) = {
get: "/api/v1/document_blocks"
};
}
rpc CreateDocumentBlock(CreateDocumentBlockRequest) returns (CreateDocumentBlockResponse) {
option (google.api.http) = {
post: "/api/v1/document_blocks"
};
}
rpc DeleteDocumentBlock(DeleteDocumentBlockRequest) returns (DeleteDocumentBlockResponse) {
option (google.api.http) = {
delete: "/api/v1/document_blocks"
};
}
rpc UpdateDocumentBlock(UpdateDocumentBlockRequest) returns (UpdateDocumentBlockResponse) {
option (google.api.http) = {
put: "/api/v1/document_blocks"
};
}
}

190
api/rag/rag.swagger.json
View File

@ -16,6 +16,143 @@
"application/json"
],
"paths": {
"/api/v1/document_blocks": {
"get": {
"operationId": "RAGService_ListDocumentBlock",
"responses": {
"200": {
"description": "A successful response.",
"schema": {
"$ref": "#/definitions/ListDocumentBlockResponse"
}
},
"default": {
"description": "An unexpected error response.",
"schema": {
"$ref": "#/definitions/rpcStatus"
}
}
},
"parameters": [
{
"name": "documentId",
"in": "query",
"required": false,
"type": "string",
"format": "int64"
}
],
"tags": [
"RAGService"
]
},
"delete": {
"operationId": "RAGService_DeleteDocumentBlock",
"responses": {
"200": {
"description": "A successful response.",
"schema": {
"$ref": "#/definitions/DeleteDocumentBlockResponse"
}
},
"default": {
"description": "An unexpected error response.",
"schema": {
"$ref": "#/definitions/rpcStatus"
}
}
},
"parameters": [
{
"name": "id",
"in": "query",
"required": false,
"type": "string",
"format": "int64"
}
],
"tags": [
"RAGService"
]
},
"post": {
"operationId": "RAGService_CreateDocumentBlock",
"responses": {
"200": {
"description": "A successful response.",
"schema": {
"$ref": "#/definitions/CreateDocumentBlockResponse"
}
},
"default": {
"description": "An unexpected error response.",
"schema": {
"$ref": "#/definitions/rpcStatus"
}
}
},
"parameters": [
{
"name": "documentId",
"in": "query",
"required": false,
"type": "string",
"format": "int64"
},
{
"name": "order",
"in": "query",
"required": false,
"type": "string",
"format": "int64"
},
{
"name": "content",
"in": "query",
"required": false,
"type": "string"
}
],
"tags": [
"RAGService"
]
},
"put": {
"operationId": "RAGService_UpdateDocumentBlock",
"responses": {
"200": {
"description": "A successful response.",
"schema": {
"$ref": "#/definitions/UpdateDocumentBlockResponse"
}
},
"default": {
"description": "An unexpected error response.",
"schema": {
"$ref": "#/definitions/rpcStatus"
}
}
},
"parameters": [
{
"name": "id",
"in": "query",
"required": false,
"type": "string",
"format": "int64"
},
{
"name": "content",
"in": "query",
"required": false,
"type": "string"
}
],
"tags": [
"RAGService"
]
}
},
"/api/v1/documents": {
"get": {
"operationId": "RAGService_ListDocument",
@ -150,6 +287,14 @@
}
},
"definitions": {
"CreateDocumentBlockResponse": {
"type": "object",
"properties": {
"documentBlock": {
"$ref": "#/definitions/DocumentBlock"
}
}
},
"CreateDocumentResponse": {
"type": "object",
"properties": {
@ -166,6 +311,9 @@
}
}
},
"DeleteDocumentBlockResponse": {
"type": "object"
},
"Document": {
"type": "object",
"properties": {
@ -194,6 +342,28 @@
}
}
},
"DocumentBlock": {
"type": "object",
"properties": {
"id": {
"type": "string",
"format": "int64"
},
"content": {
"type": "string"
},
"documentId": {
"type": "string",
"format": "int64"
},
"createdAt": {
"type": "string"
},
"updatedAt": {
"type": "string"
}
}
},
"Library": {
"type": "object",
"properties": {
@ -216,6 +386,18 @@
},
"title": "import \"google.golang.org/grpc/health/grpc_health_v1\";"
},
"ListDocumentBlockResponse": {
"type": "object",
"properties": {
"documentBlocks": {
"type": "array",
"items": {
"type": "object",
"$ref": "#/definitions/DocumentBlock"
}
}
}
},
"ListDocumentResponse": {
"type": "object",
"properties": {
@ -240,6 +422,14 @@
}
}
},
"UpdateDocumentBlockResponse": {
"type": "object",
"properties": {
"documentBlock": {
"$ref": "#/definitions/DocumentBlock"
}
}
},
"protobufAny": {
"type": "object",
"properties": {

2
consts/auth.go
View File

@ -17,4 +17,6 @@ var (
ErrNotBearerType = errors.New("不是 Bearer 类型")
ErrEmptyResponse = errors.New("我们的服务器返回了空请求，可能某些环节出了问题")
ErrTokenError = errors.New("token 类型错误")
ErrNotYourResource = errors.New("你不能修改这个资源，因为它不是你创建的。")
)

9
consts/document.go Normal file
View File

@ -0,0 +1,9 @@
package consts
import "errors"
var (
ErrDocumentNotFound = errors.New("找不到该文档。")
ErrDocumentEmpty = errors.New("文档内容为空。")
ErrDocumentInvalid = errors.New("无效的文档。")
)

7
consts/library.go Normal file
View File

@ -0,0 +1,7 @@
package consts
import "errors"
var (
ErrLibraryNotFound = errors.New("未找到指定的资料库")
)

9
consts/validate.go Normal file
View File

@ -0,0 +1,9 @@
package consts
import "errors"
var (
ErrValidateNameMustBeNotEmpty = errors.New("名称不能为空")
ErrValidateNameTooLong = errors.New("名称太长了")
ErrValidateNameTooShort = errors.New("名称太短了")
)

3
ent/schema/documentblock.go
View File

@ -16,7 +16,6 @@ type DocumentBlock struct {
// Fields of the DocumentBlock.
func (DocumentBlock) Fields() []ent.Field {
return []ent.Field{
field.String("name"),
field.String("content"),
// order
field.Int64("order"),
@ -38,6 +37,6 @@ func (DocumentBlock) Edges() []ent.Edge {
func (DocumentBlock) Indexes() []ent.Index {
return []ent.Index{
index.Fields("name", "user_id", "order"),
index.Fields("user_id", "order"),
}
}

2
internal/cmd/migrate.go
View File

@ -43,7 +43,7 @@ var migrateCommand = &cobra.Command{
var genMigrateCommand = &cobra.Command{
Use: "gen-migrate [name]",
Short: "新建 ent 迁移",
Long: "从 internal/ent 中新建迁移。在这之前，需要运行 go generate ./internal/ent",
Long: "从 internal/ent 中新建迁移。在这之前，需要运行 go generate ./ent",
Run: func(cmd *cobra.Command, args []string) {
generateMigration()
},

123
internal/logic/document.go
View File

@ -2,9 +2,11 @@ package logic
import (
"context"
"errors"
"leafdev.top/leaf/rag/consts"
"leafdev.top/leaf/rag/ent"
"leafdev.top/leaf/rag/ent/document"
"leafdev.top/leaf/rag/ent/documentblock"
"leafdev.top/leaf/rag/models"
)
type DocumentLogic struct {
@ -17,41 +19,142 @@ func NewDocumentLogic() *DocumentLogic {
}
}
func (d DocumentLogic) ListDocument(ctx context.Context, libraryId int64) ([]*ent.Document, error) {
library, err := d.LibraryLogic.ExistsByUser(ctx, int(libraryId))
func (d DocumentLogic) ListDocument(ctx context.Context, libraryId int) ([]*ent.Document, error) {
library, err := d.LibraryLogic.ExistsByUser(ctx, libraryId)
if err != nil {
return nil, err
}
if !library {
return nil, errors.New("我们在你的账户下找不到这个资料库")
return nil, consts.ErrLibraryNotFound
}
return orm.Document.Query().Where(
document.UserIDEQ(GetUserId(ctx)),
document.LibraryIDEQ(libraryId),
document.LibraryIDEQ(int64(libraryId)),
).All(ctx)
}
func (d DocumentLogic) CreateDocument(ctx context.Context, libraryId int64, name string) (*ent.Document, error) {
func (d DocumentLogic) CreateDocument(ctx context.Context, libraryId int, name string) (*ent.Document, error) {
// 检查 Library 是否存在
exists, err := d.LibraryLogic.ExistsByUser(ctx, int(libraryId))
exists, err := d.LibraryLogic.ExistsByUser(ctx, libraryId)
if err != nil {
return nil, err
}
if !exists {
return nil, errors.New("library not exists")
return nil, consts.ErrLibraryNotFound
}
if name == "" {
return nil, errors.New("name cannot be empty")
return nil, consts.ErrValidateNameMustBeNotEmpty
}
return orm.Document.Create().
SetName(name).
SetUserID(GetUserId(ctx)).
SetLibraryID(libraryId).
SetLibraryID(int64(libraryId)).
Save(ctx)
}
func (d DocumentLogic) ExistsByUser(ctx context.Context, libraryId int) (bool, error) {
return orm.Document.Query().Where(
document.UserIDEQ(GetUserId(ctx)),
document.LibraryIDEQ(int64(libraryId)),
).Exist(ctx)
}
func (d DocumentLogic) CreateDocumentBlock(ctx context.Context, documentId int, createDocumentModel *models.CreateDocumentBlock) (*ent.DocumentBlock, error) {
// 检查 Document 是否存在
exists, err := d.ExistsByUser(ctx, int(documentId))
if err != nil {
return nil, err
}
if !exists {
return nil, consts.ErrDocumentNotFound
}
return orm.DocumentBlock.Create().
SetOrder(createDocumentModel.Order).
SetDocumentID(int64(documentId)).
SetUserID(GetUserId(ctx)).
SetContent(createDocumentModel.Content).
Save(ctx)
}
func (d DocumentLogic) ListDocumentBlock(ctx context.Context, documentId int) ([]*ent.DocumentBlock, error) {
exists, err := d.ExistsByUser(ctx, documentId)
if err != nil {
return nil, err
}
if !exists {
return nil, consts.ErrDocumentNotFound
}
blocks, err := orm.DocumentBlock.Query().Where(
documentblock.DocumentID(int64(documentId)),
).All(ctx)
return blocks, err
}
func (d DocumentLogic) UpdateDocumentBlock(ctx context.Context, documentId int, blockId int, content string) (*ent.DocumentBlock, error) {
exists, err := d.ExistsByUser(ctx, documentId)
if err != nil {
return nil, err
}
if !exists {
return nil, consts.ErrDocumentNotFound
}
block, err := d.GetBlockById(ctx, blockId)
if err != nil {
return nil, err
}
if block.DocumentID != int64(documentId) {
return nil, consts.ErrDocumentNotFound
}
if block.UserID != GetUserId(ctx) {
return nil, consts.ErrNotYourResource
}
block.Content = content
return block.Update().Save(ctx)
}
func (d DocumentLogic) DeleteDocumentBlockById(ctx context.Context, documentId int, blockId int) error {
exists, err := d.ExistsByUser(ctx, documentId)
if err != nil {
return err
}
if !exists {
return consts.ErrDocumentNotFound
}
_, err = orm.DocumentBlock.Delete().Where(
documentblock.DocumentIDEQ(int64(documentId)),
documentblock.IDEQ(blockId),
).Exec(ctx)
return err
}
func (d DocumentLogic) GetBlockById(ctx context.Context, documentBlockId int) (*ent.DocumentBlock, error) {
return orm.DocumentBlock.Query().Where(documentblock.IDEQ(documentBlockId)).First(ctx)
}
func (d DocumentLogic) BlockExistsById(ctx context.Context, id int) (bool, error) {
return orm.Document.Query().Where(document.ID(id)).Exist(ctx)
}
func (d DocumentLogic) Exists(ctx context.Context, id int) (bool, error) {
return orm.Document.Query().Where(document.ID(id)).Exist(ctx)
}

4
internal/logic/library.go
View File

@ -2,7 +2,7 @@ package logic
import (
"context"
"errors"
"leafdev.top/leaf/rag/consts"
"leafdev.top/leaf/rag/ent"
"leafdev.top/leaf/rag/ent/library"
"leafdev.top/leaf/rag/internal/providers/helper"
@ -27,7 +27,7 @@ func (l *LibraryLogic) ListLibrary(ctx context.Context, page int) ([]*ent.Librar
func (l *LibraryLogic) CreateLibrary(ctx context.Context, name string) (*ent.Library, error) {
// name 不能为空
if name == "" {
return nil, errors.New("名称不能为空")
return nil, consts.ErrValidateNameMustBeNotEmpty
}
return orm.Library.Create().

46
internal/services/rag/documentService.go
View File

@ -4,10 +4,11 @@ import (
"context"
"github.com/mitchellh/mapstructure"
ragApi "leafdev.top/leaf/rag/api/rag"
"leafdev.top/leaf/rag/models"
)
func (Service) ListDocument(ctx context.Context, req *ragApi.ListDocumentRequest) (*ragApi.ListDocumentResponse, error) {
document, err := documentLogic.ListDocument(ctx, req.LibraryId)
document, err := documentLogic.ListDocument(ctx, int(req.LibraryId))
if err != nil {
return nil, err
}
@ -29,7 +30,7 @@ func (Service) ListDocument(ctx context.Context, req *ragApi.ListDocumentRequest
}
func (Service) CreateDocument(ctx context.Context, req *ragApi.CreateDocumentRequest) (*ragApi.CreateDocumentResponse, error) {
document, err := documentLogic.CreateDocument(ctx, req.LibraryId, req.Name)
document, err := documentLogic.CreateDocument(ctx, int(req.LibraryId), req.Name)
if err != nil {
return nil, err
}
@ -46,3 +47,44 @@ func (Service) CreateDocument(ctx context.Context, req *ragApi.CreateDocumentReq
//func (Service) UpdateDocument(ctx context.Context, _ *ragApi.UpdateDocumentRequest) (*ragApi.UpdateDocumentResponse, error) {
// panic("implement me")
//}
//func (Service) DeleteDocument(ctx context.Context, _ *ragApi.DeleteDocumentRequest) (*ragApi.DeleteDocumentResponse, error) {
// panic("implement me")
//}
func (Service) CreateDocumentBlock(ctx context.Context, req *ragApi.CreateDocumentBlockRequest) (*ragApi.CreateDocumentBlockResponse, error) {
var createDB = &models.CreateDocumentBlock{
Order: req.Order,
Content: req.Content,
}
documentBlock, err := documentLogic.CreateDocumentBlock(ctx, int(req.DocumentId), createDB)
if err != nil {
return nil, err
}
var documentBlockApiResponse = ragApi.CreateDocumentBlockResponse{
DocumentBlock: &ragApi.DocumentBlock{
Id: int64(documentBlock.ID),
Content: documentBlock.Content,
DocumentId: documentBlock.DocumentID,
CreatedAt: documentBlock.CreatedAt.String(),
UpdatedAt: documentBlock.UpdatedAt.String(),
},
}
return &documentBlockApiResponse, nil
}
func (Service) ListDocumentBlock(ctx context.Context, req *ragApi.ListDocumentBlockRequest) (*ragApi.ListDocumentBlockResponse, error) {
panic("implement me")
}
func (Service) UpdateDocumentBlock(ctx context.Context, req *ragApi.UpdateDocumentBlockRequest) (*ragApi.UpdateDocumentBlockResponse, error) {
panic("implement me")
}
func (Service) DeleteDocumentBlock(ctx context.Context, req *ragApi.DeleteDocumentBlockRequest) (*ragApi.DeleteDocumentBlockResponse, error) {
panic("implement me")
}

10
models/document.go Normal file
View File

@ -0,0 +1,10 @@
package models
//type DocumentModels struct {
// CreateDocumentBlock CreateDocumentBlock
//}
type CreateDocumentBlock struct {
Order int64 `json:"order"`
Content string `json:"content"`
}