fix:edit agent API

desc/wechat/agent.api

@@ -61,6 +61,187 @@ type (
         // Agent information | Agent数据
         Data AgentInfo `json:"data"`
     }
+
+	VectorModel {
+		Model *string `json:"model"`
+		Name *string `json:"name"`
+		CharsPointsPrice *uint64 `json:"charsPointsPrice"`
+		DefaultToken *uint64 `json:"defaultToken"`
+		MaxToken *uint64 `json:"maxToken"`
+		Weight *uint64 `json:"weight"`
+	}
+
+	AgentModel {
+		Model *string `json:"model"`
+		Name *string `json:"name"`
+		MaxContext *uint64 `json:"maxContext"`
+		MaxResponse *uint64 `json:"maxResponse"`
+		CharsPointsPrice *uint64 `json:"charsPointsPrice"`
+	}
+
+	DatasetId {
+		ID *string  `json:"id"`
+		ParentID *string `json:"parentId"`
+		TeamId *string `json:"teamId"`
+		TmbId *string `json:"tmbId"`
+		Type *string `json:"type"`
+		Status *string `json:"status"`
+		Name *string `json:"name"`
+		VectorModel *string `json:"vectorModel"`
+		AgentModel *string `json:"agentModel"`
+		Intro *string `json:"intro"`
+		Permission *string `json:"permission"`
+	}
+
+	// Dataset info | 知识库详情
+	DatasetInfo {
+		ID *string  `json:"id"`
+		ParentID *string `json:"parentId"`
+		TeamId *string `json:"teamId"`
+		TmbId *string `json:"tmbId"`
+		Type *string `json:"type"`
+		Name *string `json:"name"`
+		Intro *string `json:"intro"`
+		Status *string `json:"status"`
+		Avatar *string `json:"avatar"`
+		VectorModel VectorModel `json:"vectorModel"`
+		AgentModel AgentModel `json:"agentModel"`
+		Permission *string `json:"permission"`
+		CanWrite *bool `json:"canWrite"`
+		IsOwner *bool `json:"isOwner"`
+	}
+
+	Index {
+		defaultIndex *bool `json:"canWrite"`
+		Type *string `json:"type"`
+		DataId *string  `json:"dataId"`
+		Text *string `json:"text"`
+		ID *string `json:"id"`
+	}
+
+	// Collection Info | 集合详情
+	CollectionInfo {
+		ID *string  `json:"id"`
+		ParentID *string `json:"parentId"`
+		TmbId *string `json:"tmbId,optional"`
+		Type *string `json:"type"`
+		Name *string `json:"name"`
+		DataAmount *uint64 `json:"dataAmount,optional"`
+		TrainingAmount *uint64 `json:"trainingAmount,optional"`
+		TrainingType *string `json:"trainingType,optional"`
+		ChunkSize *uint64 `json:"chunkSize,optional"`
+		ChunkSplitter *string `json:"chunkSplitter,optional"`
+		QaPrompt *string `json:"qaPrompt,optional"`
+		RawTextLength *uint64 `json:"rawTextLength,optional"`
+		CanWrite *bool `json:"canWrite,optional"`
+		SourceName *string `json:"sourceName,optional"`
+		DatasetId DatasetId `json:"datasetId,optional"`
+	}
+
+	CollectionSimpleInfo {
+		ID *string  `json:"id"`
+		ParentID *string `json:"parentId"`
+		TmbId *string `json:"tmbId,optional"`
+		Type *string `json:"type"`
+		Name *string `json:"name"`
+		DataAmount *uint64 `json:"dataAmount,optional"`
+		TrainingAmount *uint64 `json:"trainingAmount,optional"`
+	}
+
+	DataInfo {
+		ID *string  `json:"id"`
+		Q *string `json:"q"`
+		A *string `json:"a"`
+		ChunkIndex *uint64 `json:"chunkIndex"`
+		Indexes []Index `json:"indexes"`
+		DatasetId *string `json:"datasetId"`
+		CollectionId *string `json:"collectionId"`
+		SourceName *string `json:"sourceName"`
+		SourceId *string `json:"sourceId"`
+		CanWrite *bool `json:"canWrite"`
+		IsOwner *bool `json:"isOwner"`
+	}
+
+	// Get collection list request params | Collection列表请求参数
+	CollectionListReq {
+		PageNum   *int    `json:"pageNum" validate:"required,number,gt=0"`
+
+		PageSize  *int    `json:"pageSize" validate:"required,number,lt=100000"`
+
+		DatasetId  *string `json:"datasetId" validate:"required"`
+	}
+
+	// Collection list response | Collection List信息返回体
+	CollectionListResp {
+		BaseDataInfo
+
+		// Agent information | Agent数据
+		Data []CollectionSimpleInfo `json:"data"`
+
+		PageNum   *int    `json:"pageNum" validate:"required,number,gt=0"`
+
+		PageSize  *int    `json:"pageSize" validate:"required,number,lt=100000"`
+
+		Total     *int    `json:"total"`
+	}
+
+	// Get collection list request params | Collection列表请求参数
+	DataListReq {
+		PageNum   *int    `json:"pageNum" validate:"required,number,gt=0"`
+
+		PageSize  *int    `json:"pageSize" validate:"required,number,lt=100000"`
+
+		CollectionId  *string `json:"collectionId" validate:"required"`
+	}
+
+	// Data list response | Data List信息返回体
+	DataListResp {
+		BaseDataInfo
+
+		// Agent information | Agent数据
+		Data []DataInfo `json:"data"`
+
+		PageNum   *int    `json:"pageNum" validate:"required,number,gt=0"`
+
+		PageSize  *int    `json:"pageSize" validate:"required,number,lt=100000"`
+
+		Total     *int    `json:"total"`
+	}
+
+	// Data create request | 信息请求体
+	CreateDataInfoReq {
+		CollectionId *string `json:"collectionId" validate:"required"`
+
+		// Q
+		 Q *string `json:"q"`
+
+		// A
+		 A *string `json:"a"`
+
+		// Indexes | 索引
+		//Indexes []IndexSingle `json:"indexes,optional"`
+	}
+
+//	IndexSingle {
+//		Text A *string `json:"text"`
+//	}
+
+	// Data create request | 信息返回体
+	UpdateDataInfoReq {
+		CollectionId *string `json:"collectionId" validate:"required"`
+
+		// ID
+		ID *string `json:"id"`
+
+		// Q
+		Q *string `json:"q"`
+
+		// A
+		A *string `json:"a"`
+
+		// Indexes | 索引
+		//Indexes []IndexSingle `json:"indexes,optional"`
+	}
 )
 
 @server(
@@ -86,7 +267,23 @@ service Wechat {
     @handler getAgentList
     post /agent/list (AgentListReq) returns (AgentListResp)
 
-    // Get agent by ID | 通过ID获取Agent
-    @handler getAgentById
-    post /agent (IDReq) returns (AgentInfoResp)
+	// Get agent by ID | 通过ID获取Agent
+	@handler getAgentById
+	post /agent (IDReq) returns (AgentInfoResp)
+
+	// Get collect list | 获取collection列表
+    @handler getAgentCollectionList
+    post /agent/collection/list (CollectionListReq) returns (CollectionListResp)
+
+	// Get data list | 获取data列表
+	@handler getAgentDataList
+	post /agent/data/list (DataListReq) returns (DataListResp)
+
+	// Create data | 添加data
+	@handler createAgentData
+	post /agent/data/create (CreateDataInfoReq) returns (BaseDataInfo)
+
+	// Update data | 修改data
+	@handler updateAgentData
+	post /agent/data/update (UpdateDataInfoReq) returns (BaseDataInfo)
 }

ent/agent.go

@@ -35,6 +35,10 @@ type Agent struct {
 	Examples string `json:"examples,omitempty"`
 	// organization_id | 租户ID
 	OrganizationID uint64 `json:"organization_id,omitempty"`
+	// dataset_id | 知识库ID
+	DatasetID string `json:"dataset_id,omitempty"`
+	// collection_id | 集合ID
+	CollectionID string `json:"collection_id,omitempty"`
 	// Edges holds the relations/edges for other nodes in the graph.
 	// The values are being populated by the AgentQuery when eager-loading is set.
 	Edges        AgentEdges `json:"edges"`
@@ -66,7 +70,7 @@ func (*Agent) scanValues(columns []string) ([]any, error) {
 		switch columns[i] {
 		case agent.FieldID, agent.FieldStatus, agent.FieldOrganizationID:
 			values[i] = new(sql.NullInt64)
-		case agent.FieldName, agent.FieldRole, agent.FieldBackground, agent.FieldExamples:
+		case agent.FieldName, agent.FieldRole, agent.FieldBackground, agent.FieldExamples, agent.FieldDatasetID, agent.FieldCollectionID:
 			values[i] = new(sql.NullString)
 		case agent.FieldCreatedAt, agent.FieldUpdatedAt, agent.FieldDeletedAt:
 			values[i] = new(sql.NullTime)
@@ -145,6 +149,18 @@ func (a *Agent) assignValues(columns []string, values []any) error {
 			} else if value.Valid {
 				a.OrganizationID = uint64(value.Int64)
 			}
+		case agent.FieldDatasetID:
+			if value, ok := values[i].(*sql.NullString); !ok {
+				return fmt.Errorf("unexpected type %T for field dataset_id", values[i])
+			} else if value.Valid {
+				a.DatasetID = value.String
+			}
+		case agent.FieldCollectionID:
+			if value, ok := values[i].(*sql.NullString); !ok {
+				return fmt.Errorf("unexpected type %T for field collection_id", values[i])
+			} else if value.Valid {
+				a.CollectionID = value.String
+			}
 		default:
 			a.selectValues.Set(columns[i], values[i])
 		}
@@ -212,6 +228,12 @@ func (a *Agent) String() string {
 	builder.WriteString(", ")
 	builder.WriteString("organization_id=")
 	builder.WriteString(fmt.Sprintf("%v", a.OrganizationID))
+	builder.WriteString(", ")
+	builder.WriteString("dataset_id=")
+	builder.WriteString(a.DatasetID)
+	builder.WriteString(", ")
+	builder.WriteString("collection_id=")
+	builder.WriteString(a.CollectionID)
 	builder.WriteByte(')')
 	return builder.String()
 }

ent/agent/agent.go

@@ -33,6 +33,10 @@ const (
 	FieldExamples = "examples"
 	// FieldOrganizationID holds the string denoting the organization_id field in the database.
 	FieldOrganizationID = "organization_id"
+	// FieldDatasetID holds the string denoting the dataset_id field in the database.
+	FieldDatasetID = "dataset_id"
+	// FieldCollectionID holds the string denoting the collection_id field in the database.
+	FieldCollectionID = "collection_id"
 	// EdgeWxAgent holds the string denoting the wx_agent edge name in mutations.
 	EdgeWxAgent = "wx_agent"
 	// Table holds the table name of the agent in the database.
@@ -58,6 +62,8 @@ var Columns = []string{
 	FieldBackground,
 	FieldExamples,
 	FieldOrganizationID,
+	FieldDatasetID,
+	FieldCollectionID,
 }
 
 // ValidColumn reports if the column name is valid (part of the table columns).
@@ -102,6 +108,10 @@ var (
 	ExamplesValidator func(string) error
 	// OrganizationIDValidator is a validator for the "organization_id" field. It is called by the builders before save.
 	OrganizationIDValidator func(uint64) error
+	// DatasetIDValidator is a validator for the "dataset_id" field. It is called by the builders before save.
+	DatasetIDValidator func(string) error
+	// CollectionIDValidator is a validator for the "collection_id" field. It is called by the builders before save.
+	CollectionIDValidator func(string) error
 )
 
 // OrderOption defines the ordering options for the Agent queries.
@@ -157,6 +167,16 @@ func ByOrganizationID(opts ...sql.OrderTermOption) OrderOption {
 	return sql.OrderByField(FieldOrganizationID, opts...).ToFunc()
 }
 
+// ByDatasetID orders the results by the dataset_id field.
+func ByDatasetID(opts ...sql.OrderTermOption) OrderOption {
+	return sql.OrderByField(FieldDatasetID, opts...).ToFunc()
+}
+
+// ByCollectionID orders the results by the collection_id field.
+func ByCollectionID(opts ...sql.OrderTermOption) OrderOption {
+	return sql.OrderByField(FieldCollectionID, opts...).ToFunc()
+}
+
 // ByWxAgentCount orders the results by wx_agent count.
 func ByWxAgentCount(opts ...sql.OrderTermOption) OrderOption {
 	return func(s *sql.Selector) {

ent/agent/where.go

@@ -100,6 +100,16 @@ func OrganizationID(v uint64) predicate.Agent {
 	return predicate.Agent(sql.FieldEQ(FieldOrganizationID, v))
 }
 
+// DatasetID applies equality check predicate on the "dataset_id" field. It's identical to DatasetIDEQ.
+func DatasetID(v string) predicate.Agent {
+	return predicate.Agent(sql.FieldEQ(FieldDatasetID, v))
+}
+
+// CollectionID applies equality check predicate on the "collection_id" field. It's identical to CollectionIDEQ.
+func CollectionID(v string) predicate.Agent {
+	return predicate.Agent(sql.FieldEQ(FieldCollectionID, v))
+}
+
 // CreatedAtEQ applies the EQ predicate on the "created_at" field.
 func CreatedAtEQ(v time.Time) predicate.Agent {
 	return predicate.Agent(sql.FieldEQ(FieldCreatedAt, v))
@@ -600,6 +610,136 @@ func OrganizationIDLTE(v uint64) predicate.Agent {
 	return predicate.Agent(sql.FieldLTE(FieldOrganizationID, v))
 }
 
+// DatasetIDEQ applies the EQ predicate on the "dataset_id" field.
+func DatasetIDEQ(v string) predicate.Agent {
+	return predicate.Agent(sql.FieldEQ(FieldDatasetID, v))
+}
+
+// DatasetIDNEQ applies the NEQ predicate on the "dataset_id" field.
+func DatasetIDNEQ(v string) predicate.Agent {
+	return predicate.Agent(sql.FieldNEQ(FieldDatasetID, v))
+}
+
+// DatasetIDIn applies the In predicate on the "dataset_id" field.
+func DatasetIDIn(vs ...string) predicate.Agent {
+	return predicate.Agent(sql.FieldIn(FieldDatasetID, vs...))
+}
+
+// DatasetIDNotIn applies the NotIn predicate on the "dataset_id" field.
+func DatasetIDNotIn(vs ...string) predicate.Agent {
+	return predicate.Agent(sql.FieldNotIn(FieldDatasetID, vs...))
+}
+
+// DatasetIDGT applies the GT predicate on the "dataset_id" field.
+func DatasetIDGT(v string) predicate.Agent {
+	return predicate.Agent(sql.FieldGT(FieldDatasetID, v))
+}
+
+// DatasetIDGTE applies the GTE predicate on the "dataset_id" field.
+func DatasetIDGTE(v string) predicate.Agent {
+	return predicate.Agent(sql.FieldGTE(FieldDatasetID, v))
+}
+
+// DatasetIDLT applies the LT predicate on the "dataset_id" field.
+func DatasetIDLT(v string) predicate.Agent {
+	return predicate.Agent(sql.FieldLT(FieldDatasetID, v))
+}
+
+// DatasetIDLTE applies the LTE predicate on the "dataset_id" field.
+func DatasetIDLTE(v string) predicate.Agent {
+	return predicate.Agent(sql.FieldLTE(FieldDatasetID, v))
+}
+
+// DatasetIDContains applies the Contains predicate on the "dataset_id" field.
+func DatasetIDContains(v string) predicate.Agent {
+	return predicate.Agent(sql.FieldContains(FieldDatasetID, v))
+}
+
+// DatasetIDHasPrefix applies the HasPrefix predicate on the "dataset_id" field.
+func DatasetIDHasPrefix(v string) predicate.Agent {
+	return predicate.Agent(sql.FieldHasPrefix(FieldDatasetID, v))
+}
+
+// DatasetIDHasSuffix applies the HasSuffix predicate on the "dataset_id" field.
+func DatasetIDHasSuffix(v string) predicate.Agent {
+	return predicate.Agent(sql.FieldHasSuffix(FieldDatasetID, v))
+}
+
+// DatasetIDEqualFold applies the EqualFold predicate on the "dataset_id" field.
+func DatasetIDEqualFold(v string) predicate.Agent {
+	return predicate.Agent(sql.FieldEqualFold(FieldDatasetID, v))
+}
+
+// DatasetIDContainsFold applies the ContainsFold predicate on the "dataset_id" field.
+func DatasetIDContainsFold(v string) predicate.Agent {
+	return predicate.Agent(sql.FieldContainsFold(FieldDatasetID, v))
+}
+
+// CollectionIDEQ applies the EQ predicate on the "collection_id" field.
+func CollectionIDEQ(v string) predicate.Agent {
+	return predicate.Agent(sql.FieldEQ(FieldCollectionID, v))
+}
+
+// CollectionIDNEQ applies the NEQ predicate on the "collection_id" field.
+func CollectionIDNEQ(v string) predicate.Agent {
+	return predicate.Agent(sql.FieldNEQ(FieldCollectionID, v))
+}
+
+// CollectionIDIn applies the In predicate on the "collection_id" field.
+func CollectionIDIn(vs ...string) predicate.Agent {
+	return predicate.Agent(sql.FieldIn(FieldCollectionID, vs...))
+}
+
+// CollectionIDNotIn applies the NotIn predicate on the "collection_id" field.
+func CollectionIDNotIn(vs ...string) predicate.Agent {
+	return predicate.Agent(sql.FieldNotIn(FieldCollectionID, vs...))
+}
+
+// CollectionIDGT applies the GT predicate on the "collection_id" field.
+func CollectionIDGT(v string) predicate.Agent {
+	return predicate.Agent(sql.FieldGT(FieldCollectionID, v))
+}
+
+// CollectionIDGTE applies the GTE predicate on the "collection_id" field.
+func CollectionIDGTE(v string) predicate.Agent {
+	return predicate.Agent(sql.FieldGTE(FieldCollectionID, v))
+}
+
+// CollectionIDLT applies the LT predicate on the "collection_id" field.
+func CollectionIDLT(v string) predicate.Agent {
+	return predicate.Agent(sql.FieldLT(FieldCollectionID, v))
+}
+
+// CollectionIDLTE applies the LTE predicate on the "collection_id" field.
+func CollectionIDLTE(v string) predicate.Agent {
+	return predicate.Agent(sql.FieldLTE(FieldCollectionID, v))
+}
+
+// CollectionIDContains applies the Contains predicate on the "collection_id" field.
+func CollectionIDContains(v string) predicate.Agent {
+	return predicate.Agent(sql.FieldContains(FieldCollectionID, v))
+}
+
+// CollectionIDHasPrefix applies the HasPrefix predicate on the "collection_id" field.
+func CollectionIDHasPrefix(v string) predicate.Agent {
+	return predicate.Agent(sql.FieldHasPrefix(FieldCollectionID, v))
+}
+
+// CollectionIDHasSuffix applies the HasSuffix predicate on the "collection_id" field.
+func CollectionIDHasSuffix(v string) predicate.Agent {
+	return predicate.Agent(sql.FieldHasSuffix(FieldCollectionID, v))
+}
+
+// CollectionIDEqualFold applies the EqualFold predicate on the "collection_id" field.
+func CollectionIDEqualFold(v string) predicate.Agent {
+	return predicate.Agent(sql.FieldEqualFold(FieldCollectionID, v))
+}
+
+// CollectionIDContainsFold applies the ContainsFold predicate on the "collection_id" field.
+func CollectionIDContainsFold(v string) predicate.Agent {
+	return predicate.Agent(sql.FieldContainsFold(FieldCollectionID, v))
+}
+
 // HasWxAgent applies the HasEdge predicate on the "wx_agent" edge.
 func HasWxAgent() predicate.Agent {
 	return predicate.Agent(func(s *sql.Selector) {

ent/agent_create.go

@@ -125,6 +125,18 @@ func (ac *AgentCreate) SetOrganizationID(u uint64) *AgentCreate {
 	return ac
 }
 
+// SetDatasetID sets the "dataset_id" field.
+func (ac *AgentCreate) SetDatasetID(s string) *AgentCreate {
+	ac.mutation.SetDatasetID(s)
+	return ac
+}
+
+// SetCollectionID sets the "collection_id" field.
+func (ac *AgentCreate) SetCollectionID(s string) *AgentCreate {
+	ac.mutation.SetCollectionID(s)
+	return ac
+}
+
 // SetID sets the "id" field.
 func (ac *AgentCreate) SetID(u uint64) *AgentCreate {
 	ac.mutation.SetID(u)
@@ -259,6 +271,22 @@ func (ac *AgentCreate) check() error {
 			return &ValidationError{Name: "organization_id", err: fmt.Errorf(`ent: validator failed for field "Agent.organization_id": %w`, err)}
 		}
 	}
+	if _, ok := ac.mutation.DatasetID(); !ok {
+		return &ValidationError{Name: "dataset_id", err: errors.New(`ent: missing required field "Agent.dataset_id"`)}
+	}
+	if v, ok := ac.mutation.DatasetID(); ok {
+		if err := agent.DatasetIDValidator(v); err != nil {
+			return &ValidationError{Name: "dataset_id", err: fmt.Errorf(`ent: validator failed for field "Agent.dataset_id": %w`, err)}
+		}
+	}
+	if _, ok := ac.mutation.CollectionID(); !ok {
+		return &ValidationError{Name: "collection_id", err: errors.New(`ent: missing required field "Agent.collection_id"`)}
+	}
+	if v, ok := ac.mutation.CollectionID(); ok {
+		if err := agent.CollectionIDValidator(v); err != nil {
+			return &ValidationError{Name: "collection_id", err: fmt.Errorf(`ent: validator failed for field "Agent.collection_id": %w`, err)}
+		}
+	}
 	return nil
 }
 
@@ -328,6 +356,14 @@ func (ac *AgentCreate) createSpec() (*Agent, *sqlgraph.CreateSpec) {
 		_spec.SetField(agent.FieldOrganizationID, field.TypeUint64, value)
 		_node.OrganizationID = value
 	}
+	if value, ok := ac.mutation.DatasetID(); ok {
+		_spec.SetField(agent.FieldDatasetID, field.TypeString, value)
+		_node.DatasetID = value
+	}
+	if value, ok := ac.mutation.CollectionID(); ok {
+		_spec.SetField(agent.FieldCollectionID, field.TypeString, value)
+		_node.CollectionID = value
+	}
 	if nodes := ac.mutation.WxAgentIDs(); len(nodes) > 0 {
 		edge := &sqlgraph.EdgeSpec{
 			Rel:     sqlgraph.O2M,
@@ -528,6 +564,30 @@ func (u *AgentUpsert) AddOrganizationID(v uint64) *AgentUpsert {
 	return u
 }
 
+// SetDatasetID sets the "dataset_id" field.
+func (u *AgentUpsert) SetDatasetID(v string) *AgentUpsert {
+	u.Set(agent.FieldDatasetID, v)
+	return u
+}
+
+// UpdateDatasetID sets the "dataset_id" field to the value that was provided on create.
+func (u *AgentUpsert) UpdateDatasetID() *AgentUpsert {
+	u.SetExcluded(agent.FieldDatasetID)
+	return u
+}
+
+// SetCollectionID sets the "collection_id" field.
+func (u *AgentUpsert) SetCollectionID(v string) *AgentUpsert {
+	u.Set(agent.FieldCollectionID, v)
+	return u
+}
+
+// UpdateCollectionID sets the "collection_id" field to the value that was provided on create.
+func (u *AgentUpsert) UpdateCollectionID() *AgentUpsert {
+	u.SetExcluded(agent.FieldCollectionID)
+	return u
+}
+
 // UpdateNewValues updates the mutable fields using the new values that were set on create except the ID field.
 // Using this option is equivalent to using:
 //
@@ -733,6 +793,34 @@ func (u *AgentUpsertOne) UpdateOrganizationID() *AgentUpsertOne {
 	})
 }
 
+// SetDatasetID sets the "dataset_id" field.
+func (u *AgentUpsertOne) SetDatasetID(v string) *AgentUpsertOne {
+	return u.Update(func(s *AgentUpsert) {
+		s.SetDatasetID(v)
+	})
+}
+
+// UpdateDatasetID sets the "dataset_id" field to the value that was provided on create.
+func (u *AgentUpsertOne) UpdateDatasetID() *AgentUpsertOne {
+	return u.Update(func(s *AgentUpsert) {
+		s.UpdateDatasetID()
+	})
+}
+
+// SetCollectionID sets the "collection_id" field.
+func (u *AgentUpsertOne) SetCollectionID(v string) *AgentUpsertOne {
+	return u.Update(func(s *AgentUpsert) {
+		s.SetCollectionID(v)
+	})
+}
+
+// UpdateCollectionID sets the "collection_id" field to the value that was provided on create.
+func (u *AgentUpsertOne) UpdateCollectionID() *AgentUpsertOne {
+	return u.Update(func(s *AgentUpsert) {
+		s.UpdateCollectionID()
+	})
+}
+
 // Exec executes the query.
 func (u *AgentUpsertOne) Exec(ctx context.Context) error {
 	if len(u.create.conflict) == 0 {
@@ -1104,6 +1192,34 @@ func (u *AgentUpsertBulk) UpdateOrganizationID() *AgentUpsertBulk {
 	})
 }
 
+// SetDatasetID sets the "dataset_id" field.
+func (u *AgentUpsertBulk) SetDatasetID(v string) *AgentUpsertBulk {
+	return u.Update(func(s *AgentUpsert) {
+		s.SetDatasetID(v)
+	})
+}
+
+// UpdateDatasetID sets the "dataset_id" field to the value that was provided on create.
+func (u *AgentUpsertBulk) UpdateDatasetID() *AgentUpsertBulk {
+	return u.Update(func(s *AgentUpsert) {
+		s.UpdateDatasetID()
+	})
+}
+
+// SetCollectionID sets the "collection_id" field.
+func (u *AgentUpsertBulk) SetCollectionID(v string) *AgentUpsertBulk {
+	return u.Update(func(s *AgentUpsert) {
+		s.SetCollectionID(v)
+	})
+}
+
+// UpdateCollectionID sets the "collection_id" field to the value that was provided on create.
+func (u *AgentUpsertBulk) UpdateCollectionID() *AgentUpsertBulk {
+	return u.Update(func(s *AgentUpsert) {
+		s.UpdateCollectionID()
+	})
+}
+
 // Exec executes the query.
 func (u *AgentUpsertBulk) Exec(ctx context.Context) error {
 	if u.create.err != nil {

ent/agent_update.go

@@ -171,6 +171,34 @@ func (au *AgentUpdate) AddOrganizationID(u int64) *AgentUpdate {
 	return au
 }
 
+// SetDatasetID sets the "dataset_id" field.
+func (au *AgentUpdate) SetDatasetID(s string) *AgentUpdate {
+	au.mutation.SetDatasetID(s)
+	return au
+}
+
+// SetNillableDatasetID sets the "dataset_id" field if the given value is not nil.
+func (au *AgentUpdate) SetNillableDatasetID(s *string) *AgentUpdate {
+	if s != nil {
+		au.SetDatasetID(*s)
+	}
+	return au
+}
+
+// SetCollectionID sets the "collection_id" field.
+func (au *AgentUpdate) SetCollectionID(s string) *AgentUpdate {
+	au.mutation.SetCollectionID(s)
+	return au
+}
+
+// SetNillableCollectionID sets the "collection_id" field if the given value is not nil.
+func (au *AgentUpdate) SetNillableCollectionID(s *string) *AgentUpdate {
+	if s != nil {
+		au.SetCollectionID(*s)
+	}
+	return au
+}
+
 // AddWxAgentIDs adds the "wx_agent" edge to the Wx entity by IDs.
 func (au *AgentUpdate) AddWxAgentIDs(ids ...uint64) *AgentUpdate {
 	au.mutation.AddWxAgentIDs(ids...)
@@ -286,6 +314,16 @@ func (au *AgentUpdate) check() error {
 			return &ValidationError{Name: "organization_id", err: fmt.Errorf(`ent: validator failed for field "Agent.organization_id": %w`, err)}
 		}
 	}
+	if v, ok := au.mutation.DatasetID(); ok {
+		if err := agent.DatasetIDValidator(v); err != nil {
+			return &ValidationError{Name: "dataset_id", err: fmt.Errorf(`ent: validator failed for field "Agent.dataset_id": %w`, err)}
+		}
+	}
+	if v, ok := au.mutation.CollectionID(); ok {
+		if err := agent.CollectionIDValidator(v); err != nil {
+			return &ValidationError{Name: "collection_id", err: fmt.Errorf(`ent: validator failed for field "Agent.collection_id": %w`, err)}
+		}
+	}
 	return nil
 }
 
@@ -343,6 +381,12 @@ func (au *AgentUpdate) sqlSave(ctx context.Context) (n int, err error) {
 	if value, ok := au.mutation.AddedOrganizationID(); ok {
 		_spec.AddField(agent.FieldOrganizationID, field.TypeUint64, value)
 	}
+	if value, ok := au.mutation.DatasetID(); ok {
+		_spec.SetField(agent.FieldDatasetID, field.TypeString, value)
+	}
+	if value, ok := au.mutation.CollectionID(); ok {
+		_spec.SetField(agent.FieldCollectionID, field.TypeString, value)
+	}
 	if au.mutation.WxAgentCleared() {
 		edge := &sqlgraph.EdgeSpec{
 			Rel:     sqlgraph.O2M,
@@ -550,6 +594,34 @@ func (auo *AgentUpdateOne) AddOrganizationID(u int64) *AgentUpdateOne {
 	return auo
 }
 
+// SetDatasetID sets the "dataset_id" field.
+func (auo *AgentUpdateOne) SetDatasetID(s string) *AgentUpdateOne {
+	auo.mutation.SetDatasetID(s)
+	return auo
+}
+
+// SetNillableDatasetID sets the "dataset_id" field if the given value is not nil.
+func (auo *AgentUpdateOne) SetNillableDatasetID(s *string) *AgentUpdateOne {
+	if s != nil {
+		auo.SetDatasetID(*s)
+	}
+	return auo
+}
+
+// SetCollectionID sets the "collection_id" field.
+func (auo *AgentUpdateOne) SetCollectionID(s string) *AgentUpdateOne {
+	auo.mutation.SetCollectionID(s)
+	return auo
+}
+
+// SetNillableCollectionID sets the "collection_id" field if the given value is not nil.
+func (auo *AgentUpdateOne) SetNillableCollectionID(s *string) *AgentUpdateOne {
+	if s != nil {
+		auo.SetCollectionID(*s)
+	}
+	return auo
+}
+
 // AddWxAgentIDs adds the "wx_agent" edge to the Wx entity by IDs.
 func (auo *AgentUpdateOne) AddWxAgentIDs(ids ...uint64) *AgentUpdateOne {
 	auo.mutation.AddWxAgentIDs(ids...)
@@ -678,6 +750,16 @@ func (auo *AgentUpdateOne) check() error {
 			return &ValidationError{Name: "organization_id", err: fmt.Errorf(`ent: validator failed for field "Agent.organization_id": %w`, err)}
 		}
 	}
+	if v, ok := auo.mutation.DatasetID(); ok {
+		if err := agent.DatasetIDValidator(v); err != nil {
+			return &ValidationError{Name: "dataset_id", err: fmt.Errorf(`ent: validator failed for field "Agent.dataset_id": %w`, err)}
+		}
+	}
+	if v, ok := auo.mutation.CollectionID(); ok {
+		if err := agent.CollectionIDValidator(v); err != nil {
+			return &ValidationError{Name: "collection_id", err: fmt.Errorf(`ent: validator failed for field "Agent.collection_id": %w`, err)}
+		}
+	}
 	return nil
 }
 
@@ -752,6 +834,12 @@ func (auo *AgentUpdateOne) sqlSave(ctx context.Context) (_node *Agent, err error
 	if value, ok := auo.mutation.AddedOrganizationID(); ok {
 		_spec.AddField(agent.FieldOrganizationID, field.TypeUint64, value)
 	}
+	if value, ok := auo.mutation.DatasetID(); ok {
+		_spec.SetField(agent.FieldDatasetID, field.TypeString, value)
+	}
+	if value, ok := auo.mutation.CollectionID(); ok {
+		_spec.SetField(agent.FieldCollectionID, field.TypeString, value)
+	}
 	if auo.mutation.WxAgentCleared() {
 		edge := &sqlgraph.EdgeSpec{
 			Rel:     sqlgraph.O2M,

ent/migrate/schema.go

@@ -21,6 +21,8 @@ var (
 		{Name: "background", Type: field.TypeString, Nullable: true, Size: 1000, Comment: "background | 背景介绍", Default: ""},
 		{Name: "examples", Type: field.TypeString, Nullable: true, Size: 5000, Comment: "examples | 对话案例", Default: ""},
 		{Name: "organization_id", Type: field.TypeUint64, Comment: "organization_id | 租户ID"},
+		{Name: "dataset_id", Type: field.TypeString, Size: 255, Comment: "dataset_id | 知识库ID"},
+		{Name: "collection_id", Type: field.TypeString, Size: 255, Comment: "collection_id | 集合ID"},
 	}
 	// AgentTable holds the schema information for the "agent" table.
 	AgentTable = &schema.Table{

ent/mutation.go

@@ -81,6 +81,8 @@ type AgentMutation struct {
 	examples           *string
 	organization_id    *uint64
 	addorganization_id *int64
+	dataset_id         *string
+	collection_id      *string
 	clearedFields      map[string]struct{}
 	wx_agent           map[uint64]struct{}
 	removedwx_agent    map[uint64]struct{}
@@ -611,6 +613,78 @@ func (m *AgentMutation) ResetOrganizationID() {
 	m.addorganization_id = nil
 }
 
+// SetDatasetID sets the "dataset_id" field.
+func (m *AgentMutation) SetDatasetID(s string) {
+	m.dataset_id = &s
+}
+
+// DatasetID returns the value of the "dataset_id" field in the mutation.
+func (m *AgentMutation) DatasetID() (r string, exists bool) {
+	v := m.dataset_id
+	if v == nil {
+		return
+	}
+	return *v, true
+}
+
+// OldDatasetID returns the old "dataset_id" field's value of the Agent entity.
+// If the Agent object wasn't provided to the builder, the object is fetched from the database.
+// An error is returned if the mutation operation is not UpdateOne, or the database query fails.
+func (m *AgentMutation) OldDatasetID(ctx context.Context) (v string, err error) {
+	if !m.op.Is(OpUpdateOne) {
+		return v, errors.New("OldDatasetID is only allowed on UpdateOne operations")
+	}
+	if m.id == nil || m.oldValue == nil {
+		return v, errors.New("OldDatasetID requires an ID field in the mutation")
+	}
+	oldValue, err := m.oldValue(ctx)
+	if err != nil {
+		return v, fmt.Errorf("querying old value for OldDatasetID: %w", err)
+	}
+	return oldValue.DatasetID, nil
+}
+
+// ResetDatasetID resets all changes to the "dataset_id" field.
+func (m *AgentMutation) ResetDatasetID() {
+	m.dataset_id = nil
+}
+
+// SetCollectionID sets the "collection_id" field.
+func (m *AgentMutation) SetCollectionID(s string) {
+	m.collection_id = &s
+}
+
+// CollectionID returns the value of the "collection_id" field in the mutation.
+func (m *AgentMutation) CollectionID() (r string, exists bool) {
+	v := m.collection_id
+	if v == nil {
+		return
+	}
+	return *v, true
+}
+
+// OldCollectionID returns the old "collection_id" field's value of the Agent entity.
+// If the Agent object wasn't provided to the builder, the object is fetched from the database.
+// An error is returned if the mutation operation is not UpdateOne, or the database query fails.
+func (m *AgentMutation) OldCollectionID(ctx context.Context) (v string, err error) {
+	if !m.op.Is(OpUpdateOne) {
+		return v, errors.New("OldCollectionID is only allowed on UpdateOne operations")
+	}
+	if m.id == nil || m.oldValue == nil {
+		return v, errors.New("OldCollectionID requires an ID field in the mutation")
+	}
+	oldValue, err := m.oldValue(ctx)
+	if err != nil {
+		return v, fmt.Errorf("querying old value for OldCollectionID: %w", err)
+	}
+	return oldValue.CollectionID, nil
+}
+
+// ResetCollectionID resets all changes to the "collection_id" field.
+func (m *AgentMutation) ResetCollectionID() {
+	m.collection_id = nil
+}
+
 // AddWxAgentIDs adds the "wx_agent" edge to the Wx entity by ids.
 func (m *AgentMutation) AddWxAgentIDs(ids ...uint64) {
 	if m.wx_agent == nil {
@@ -699,7 +773,7 @@ func (m *AgentMutation) Type() string {
 // order to get all numeric fields that were incremented/decremented, call
 // AddedFields().
 func (m *AgentMutation) Fields() []string {
-	fields := make([]string, 0, 9)
+	fields := make([]string, 0, 11)
 	if m.created_at != nil {
 		fields = append(fields, agent.FieldCreatedAt)
 	}
@@ -727,6 +801,12 @@ func (m *AgentMutation) Fields() []string {
 	if m.organization_id != nil {
 		fields = append(fields, agent.FieldOrganizationID)
 	}
+	if m.dataset_id != nil {
+		fields = append(fields, agent.FieldDatasetID)
+	}
+	if m.collection_id != nil {
+		fields = append(fields, agent.FieldCollectionID)
+	}
 	return fields
 }
 
@@ -753,6 +833,10 @@ func (m *AgentMutation) Field(name string) (ent.Value, bool) {
 		return m.Examples()
 	case agent.FieldOrganizationID:
 		return m.OrganizationID()
+	case agent.FieldDatasetID:
+		return m.DatasetID()
+	case agent.FieldCollectionID:
+		return m.CollectionID()
 	}
 	return nil, false
 }
@@ -780,6 +864,10 @@ func (m *AgentMutation) OldField(ctx context.Context, name string) (ent.Value, e
 		return m.OldExamples(ctx)
 	case agent.FieldOrganizationID:
 		return m.OldOrganizationID(ctx)
+	case agent.FieldDatasetID:
+		return m.OldDatasetID(ctx)
+	case agent.FieldCollectionID:
+		return m.OldCollectionID(ctx)
 	}
 	return nil, fmt.Errorf("unknown Agent field %s", name)
 }
@@ -852,6 +940,20 @@ func (m *AgentMutation) SetField(name string, value ent.Value) error {
 		}
 		m.SetOrganizationID(v)
 		return nil
+	case agent.FieldDatasetID:
+		v, ok := value.(string)
+		if !ok {
+			return fmt.Errorf("unexpected type %T for field %s", value, name)
+		}
+		m.SetDatasetID(v)
+		return nil
+	case agent.FieldCollectionID:
+		v, ok := value.(string)
+		if !ok {
+			return fmt.Errorf("unexpected type %T for field %s", value, name)
+		}
+		m.SetCollectionID(v)
+		return nil
 	}
 	return fmt.Errorf("unknown Agent field %s", name)
 }
@@ -982,6 +1084,12 @@ func (m *AgentMutation) ResetField(name string) error {
 	case agent.FieldOrganizationID:
 		m.ResetOrganizationID()
 		return nil
+	case agent.FieldDatasetID:
+		m.ResetDatasetID()
+		return nil
+	case agent.FieldCollectionID:
+		m.ResetCollectionID()
+		return nil
 	}
 	return fmt.Errorf("unknown Agent field %s", name)
 }

ent/runtime/runtime.go

@@ -79,6 +79,14 @@ func init() {
 	agentDescOrganizationID := agentFields[5].Descriptor()
 	// agent.OrganizationIDValidator is a validator for the "organization_id" field. It is called by the builders before save.
 	agent.OrganizationIDValidator = agentDescOrganizationID.Validators[0].(func(uint64) error)
+	// agentDescDatasetID is the schema descriptor for dataset_id field.
+	agentDescDatasetID := agentFields[6].Descriptor()
+	// agent.DatasetIDValidator is a validator for the "dataset_id" field. It is called by the builders before save.
+	agent.DatasetIDValidator = agentDescDatasetID.Validators[0].(func(string) error)
+	// agentDescCollectionID is the schema descriptor for collection_id field.
+	agentDescCollectionID := agentFields[7].Descriptor()
+	// agent.CollectionIDValidator is a validator for the "collection_id" field. It is called by the builders before save.
+	agent.CollectionIDValidator = agentDescCollectionID.Validators[0].(func(string) error)
 	batchmsgMixin := schema.BatchMsg{}.Mixin()
 	batchmsgMixinHooks1 := batchmsgMixin[1].Hooks()
 	batchmsg.Hooks[0] = batchmsgMixinHooks1[0]

ent/schema/agent.go

@@ -24,6 +24,8 @@ func (Agent) Fields() []ent.Field {
 		field.String("background").Optional().Default("").MaxLen(1000).Comment("background | 背景介绍"),
 		field.String("examples").Optional().Default("").MaxLen(5000).Comment("examples | 对话案例"),
 		field.Uint64("organization_id").Positive().Comment("organization_id | 租户ID"),
+		field.String("dataset_id").MaxLen(255).Comment("dataset_id | 知识库ID"),
+		field.String("collection_id").MaxLen(255).Comment("collection_id | 集合ID"),
 	}
 }
 

ent/set_not_nil.go

@@ -200,6 +200,54 @@ func (a *AgentCreate) SetNotNilOrganizationID(value *uint64) *AgentCreate {
 }
 
 // set field if value's pointer is not nil.
+func (a *AgentUpdate) SetNotNilDatasetID(value *string) *AgentUpdate {
+	if value != nil {
+		return a.SetDatasetID(*value)
+	}
+	return a
+}
+
+// set field if value's pointer is not nil.
+func (a *AgentUpdateOne) SetNotNilDatasetID(value *string) *AgentUpdateOne {
+	if value != nil {
+		return a.SetDatasetID(*value)
+	}
+	return a
+}
+
+// set field if value's pointer is not nil.
+func (a *AgentCreate) SetNotNilDatasetID(value *string) *AgentCreate {
+	if value != nil {
+		return a.SetDatasetID(*value)
+	}
+	return a
+}
+
+// set field if value's pointer is not nil.
+func (a *AgentUpdate) SetNotNilCollectionID(value *string) *AgentUpdate {
+	if value != nil {
+		return a.SetCollectionID(*value)
+	}
+	return a
+}
+
+// set field if value's pointer is not nil.
+func (a *AgentUpdateOne) SetNotNilCollectionID(value *string) *AgentUpdateOne {
+	if value != nil {
+		return a.SetCollectionID(*value)
+	}
+	return a
+}
+
+// set field if value's pointer is not nil.
+func (a *AgentCreate) SetNotNilCollectionID(value *string) *AgentCreate {
+	if value != nil {
+		return a.SetCollectionID(*value)
+	}
+	return a
+}
+
+// set field if value's pointer is not nil.
 func (bm *BatchMsgUpdate) SetNotNilUpdatedAt(value *time.Time) *BatchMsgUpdate {
 	if value != nil {
 		return bm.SetUpdatedAt(*value)

hook/fastgpt/collection.go

@@ -29,7 +29,7 @@ type CreateTextCollectionReq struct {
 	Text          string      `json:"text"`
 	Metadata      interface{} `json:"metadata,optional"`
 	TrainingType  string      `json:"trainingType,optional"`
-	ChunkSize     int         `json:"chunkSize,optional"`
+	ChunkSize     uint64      `json:"chunkSize,optional"`
 	ChunkSplitter string      `json:"chunkSplitter,optional"`
 	QaPrompt      string      `json:"qaPrompt,optional"`
 }
@@ -57,7 +57,7 @@ type CreateLinkCollectionReq struct {
 
 	Metadata      interface{} `json:"metadata,optional"`
 	TrainingType  string      `json:"trainingType,optional"`
-	ChunkSize     int         `json:"chunkSize,optional"`
+	ChunkSize     uint64      `json:"chunkSize,optional"`
 	ChunkSplitter string      `json:"chunkSplitter,optional"`
 	QaPrompt      string      `json:"qaPrompt,optional"`
 }
@@ -77,7 +77,7 @@ type CreateFileCollectionReq struct {
 
 	Metadata      interface{} `json:"metadata,optional"`
 	TrainingType  string      `json:"trainingType,optional"`
-	ChunkSize     int         `json:"chunkSize,optional"`
+	ChunkSize     uint64      `json:"chunkSize,optional"`
 	ChunkSplitter string      `json:"chunkSplitter,optional"`
 	QaPrompt      string      `json:"qaPrompt,optional"`
 }
@@ -100,7 +100,7 @@ type CreateFileCollectionResp struct {
 type GetCollectionListReq struct {
 	PageNum    int    `json:"pageNum,optional"`
 	PageSize   int    `json:"pageSize,optional"`
-	DatasetId  string `json:"datasetId,optional"`
+	DatasetId  string `json:"datasetId"`
 	ParentId   string `json:"parentId,optional"`
 	SearchText string `json:"searchText,optional"`
 }
@@ -144,24 +144,24 @@ type CollectionListResp struct {
 	StatusText string `json:"statusText"`
 	Message    string `json:"message"`
 	Data       struct {
-		PageNum  int `json:"pageNum"`
-		PageSize int `json:"pageSize"`
-		Data     []struct {
-			ID             string    `json:"_id"`
-			ParentID       string    `json:"parentId"`
-			TmbID          string    `json:"tmbId"`
-			Type           string    `json:"type"`
-			Name           string    `json:"name"`
-			UpdateTime     time.Time `json:"updateTime"`
-			DataAmount     int       `json:"dataAmount"`
-			TrainingAmount int       `json:"trainingAmount"`
-			CanWrite       bool      `json:"canWrite"`
-			RawLink        string    `json:"rawLink,omitempty"`
-		} `json:"data"`
-		Total int `json:"total"`
+		PageNum  int                    `json:"pageNum"`
+		PageSize int                    `json:"pageSize"`
+		Data     []CollectionSimpleInfo `json:"data"`
+		Total    int                    `json:"total"`
 	} `json:"data"`
 }
 
+type CollectionSimpleInfo struct {
+	ID             string    `json:"_id"`
+	ParentID       string    `json:"parentId"`
+	TmbID          string    `json:"tmbId"`
+	Type           string    `json:"type"`
+	Name           string    `json:"name"`
+	UpdateTime     time.Time `json:"updateTime"`
+	DataAmount     uint64    `json:"dataAmount"`
+	TrainingAmount uint64    `json:"trainingAmount"`
+}
+
 type CollectionDetailResp struct {
 	Code       int            `json:"code"`
 	StatusText string         `json:"statusText"`

hook/fastgpt/data.go

@@ -35,10 +35,10 @@ type CreateBulkDataResp struct {
 	Code       int    `json:"code"`
 	StatusText string `json:"statusText"`
 	Data       struct {
-		InsertLen int   `json:"insertLen"`
-		OverToken []any `json:"overToken"`
-		Repeat    []any `json:"repeat"`
-		Error     []any `json:"error"`
+		InsertLen uint64 `json:"insertLen"`
+		OverToken []any  `json:"overToken"`
+		Repeat    []any  `json:"repeat"`
+		Error     []any  `json:"error"`
 	} `json:"data"`
 }
 
@@ -54,7 +54,7 @@ type DataInfo struct {
 	TeamID     string `json:"teamId"`
 	Q          string `json:"q"`
 	A          string `json:"a"`
-	ChunkIndex int    `json:"chunkIndex"`
+	ChunkIndex uint64 `json:"chunkIndex"`
 	Indexes    []struct {
 		DefaultIndex bool   `json:"defaultIndex"`
 		DataID       string `json:"dataId"`
@@ -82,9 +82,9 @@ type DataListResp struct {
 			CollectionID string `json:"collectionId"`
 			Q            string `json:"q"`
 			A            string `json:"a"`
-			ChunkIndex   int    `json:"chunkIndex"`
+			ChunkIndex   uint64 `json:"chunkIndex"`
 		} `json:"data,optional"`
-		Total int `json:"total"`
+		Total uint64 `json:"total"`
 	} `json:"data,optional"`
 }
 

hook/fastgpt/search.go

@@ -5,12 +5,12 @@ import (
 )
 
 type SearchReq struct {
-	DatasetID   string `json:"datasetId"`           //知识库ID
-	Text        string `json:"text"`                //需要测试的文本
-	Limit       int    `json:"limit"`               //最大 tokens 数量
-	Similarity  int    `json:"similarity,optional"` //最低相关度（0~1，可选）
-	SearchMode  string `json:"searchMode"`          //搜索模式：embedding | fullTextRecall | mixedRecall
-	UsingReRank bool   `json:"usingReRank"`         //使用重排
+	DatasetID   string  `json:"datasetId"`           //知识库ID
+	Text        string  `json:"text"`                //需要测试的文本
+	Limit       uint64  `json:"limit"`               //最大 tokens 数量
+	Similarity  float64 `json:"similarity,optional"` //最低相关度（0~1，可选）
+	SearchMode  string  `json:"searchMode"`          //搜索模式：embedding | fullTextRecall | mixedRecall
+	UsingReRank bool    `json:"usingReRank"`         //使用重排
 }
 
 type SearchResp struct {
@@ -29,13 +29,13 @@ type SearchResp struct {
 			Score        []struct {
 				Type  string  `json:"type"`
 				Value float64 `json:"value"`
-				Index int     `json:"index"`
+				Index uint64  `json:"index"`
 			} `json:"score"`
 		} `json:"list,optional"`
 		Duration              string `json:"duration"`
 		SearchMode            string `json:"searchMode"`
-		Limit                 int    `json:"limit"`
-		Similarity            int    `json:"similarity"`
+		Limit                 uint64 `json:"limit"`
+		Similarity            uint64 `json:"similarity"`
 		UsingReRank           bool   `json:"usingReRank"`
 		UsingSimilarityFilter bool   `json:"usingSimilarityFilter"`
 	} `json:"data,optional"`

internal/handler/agent/create_agent_data_handler.go

@@ -0,0 +1,44 @@
+package agent
+
+import (
+	"net/http"
+
+	"github.com/zeromicro/go-zero/rest/httpx"
+
+	"wechat-api/internal/logic/agent"
+	"wechat-api/internal/svc"
+	"wechat-api/internal/types"
+)
+
+// swagger:route post /agent/data/create agent CreateAgentData
+//
+// Create data | 添加data
+//
+// Create data | 添加data
+//
+// Parameters:
+//  + name: body
+//    require: true
+//    in: body
+//    type: CreateDataInfoReq
+//
+// Responses:
+//  200: BaseDataInfo
+
+func CreateAgentDataHandler(svcCtx *svc.ServiceContext) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		var req types.CreateDataInfoReq
+		if err := httpx.Parse(r, &req, true); err != nil {
+			httpx.ErrorCtx(r.Context(), w, err)
+			return
+		}
+
+		l := agent.NewCreateAgentDataLogic(r.Context(), svcCtx)
+		resp, err := l.CreateAgentData(&req)
+		if err != nil {
+			httpx.ErrorCtx(r.Context(), w, err)
+		} else {
+			httpx.OkJsonCtx(r.Context(), w, resp)
+		}
+	}
+}

internal/handler/agent/get_agent_collection_list_handler.go

@@ -0,0 +1,44 @@
+package agent
+
+import (
+	"net/http"
+
+	"github.com/zeromicro/go-zero/rest/httpx"
+
+	"wechat-api/internal/logic/agent"
+	"wechat-api/internal/svc"
+	"wechat-api/internal/types"
+)
+
+// swagger:route post /agent/collection/list agent GetAgentCollectionList
+//
+// Get collect list | 获取collection列表
+//
+// Get collect list | 获取collection列表
+//
+// Parameters:
+//  + name: body
+//    require: true
+//    in: body
+//    type: CollectionListReq
+//
+// Responses:
+//  200: CollectionListResp
+
+func GetAgentCollectionListHandler(svcCtx *svc.ServiceContext) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		var req types.CollectionListReq
+		if err := httpx.Parse(r, &req, true); err != nil {
+			httpx.ErrorCtx(r.Context(), w, err)
+			return
+		}
+
+		l := agent.NewGetAgentCollectionListLogic(r.Context(), svcCtx)
+		resp, err := l.GetAgentCollectionList(&req)
+		if err != nil {
+			httpx.ErrorCtx(r.Context(), w, err)
+		} else {
+			httpx.OkJsonCtx(r.Context(), w, resp)
+		}
+	}
+}

internal/handler/agent/get_agent_data_list_handler.go

@@ -0,0 +1,44 @@
+package agent
+
+import (
+	"net/http"
+
+	"github.com/zeromicro/go-zero/rest/httpx"
+
+	"wechat-api/internal/logic/agent"
+	"wechat-api/internal/svc"
+	"wechat-api/internal/types"
+)
+
+// swagger:route post /agent/data/list agent GetAgentDataList
+//
+// Get data list | 获取data列表
+//
+// Get data list | 获取data列表
+//
+// Parameters:
+//  + name: body
+//    require: true
+//    in: body
+//    type: DataListReq
+//
+// Responses:
+//  200: DataListResp
+
+func GetAgentDataListHandler(svcCtx *svc.ServiceContext) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		var req types.DataListReq
+		if err := httpx.Parse(r, &req, true); err != nil {
+			httpx.ErrorCtx(r.Context(), w, err)
+			return
+		}
+
+		l := agent.NewGetAgentDataListLogic(r.Context(), svcCtx)
+		resp, err := l.GetAgentDataList(&req)
+		if err != nil {
+			httpx.ErrorCtx(r.Context(), w, err)
+		} else {
+			httpx.OkJsonCtx(r.Context(), w, resp)
+		}
+	}
+}

internal/handler/agent/update_agent_data_handler.go

@@ -0,0 +1,44 @@
+package agent
+
+import (
+	"net/http"
+
+	"github.com/zeromicro/go-zero/rest/httpx"
+
+	"wechat-api/internal/logic/agent"
+	"wechat-api/internal/svc"
+	"wechat-api/internal/types"
+)
+
+// swagger:route post /agent/data/update agent UpdateAgentData
+//
+// Update data | 修改data
+//
+// Update data | 修改data
+//
+// Parameters:
+//  + name: body
+//    require: true
+//    in: body
+//    type: UpdateDataInfoReq
+//
+// Responses:
+//  200: BaseDataInfo
+
+func UpdateAgentDataHandler(svcCtx *svc.ServiceContext) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		var req types.UpdateDataInfoReq
+		if err := httpx.Parse(r, &req, true); err != nil {
+			httpx.ErrorCtx(r.Context(), w, err)
+			return
+		}
+
+		l := agent.NewUpdateAgentDataLogic(r.Context(), svcCtx)
+		resp, err := l.UpdateAgentData(&req)
+		if err != nil {
+			httpx.ErrorCtx(r.Context(), w, err)
+		} else {
+			httpx.OkJsonCtx(r.Context(), w, resp)
+		}
+	}
+}

internal/handler/routes.go

@@ -146,6 +146,26 @@ func RegisterHandlers(server *rest.Server, serverCtx *svc.ServiceContext) {
 					Path:    "/agent",
 					Handler: agent.GetAgentByIdHandler(serverCtx),
 				},
+				{
+					Method:  http.MethodPost,
+					Path:    "/agent/collection/list",
+					Handler: agent.GetAgentCollectionListHandler(serverCtx),
+				},
+				{
+					Method:  http.MethodPost,
+					Path:    "/agent/data/list",
+					Handler: agent.GetAgentDataListHandler(serverCtx),
+				},
+				{
+					Method:  http.MethodPost,
+					Path:    "/agent/data/create",
+					Handler: agent.CreateAgentDataHandler(serverCtx),
+				},
+				{
+					Method:  http.MethodPost,
+					Path:    "/agent/data/update",
+					Handler: agent.UpdateAgentDataHandler(serverCtx),
+				},
 			}...,
 		),
 		rest.WithJwt(serverCtx.Config.Auth.AccessSecret),

internal/logic/agent/create_agent_data_logic.go

@@ -0,0 +1,54 @@
+package agent
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"github.com/suyuan32/simple-admin-common/msg/errormsg"
+	"wechat-api/hook/fastgpt"
+
+	"wechat-api/internal/svc"
+	"wechat-api/internal/types"
+
+	"github.com/zeromicro/go-zero/core/logx"
+)
+
+type CreateAgentDataLogic struct {
+	logx.Logger
+	ctx    context.Context
+	svcCtx *svc.ServiceContext
+}
+
+func NewCreateAgentDataLogic(ctx context.Context, svcCtx *svc.ServiceContext) *CreateAgentDataLogic {
+	return &CreateAgentDataLogic{
+		Logger: logx.WithContext(ctx),
+		ctx:    ctx,
+		svcCtx: svcCtx}
+}
+
+func (l *CreateAgentDataLogic) CreateAgentData(req *types.CreateDataInfoReq) (*types.BaseDataInfo, error) {
+	params := fastgpt.CreateBulkDataReq{}
+	params.CollectionID = *req.CollectionId
+	params.TrainingMode = "chunk"
+	pair := make([]fastgpt.DataQuestion, 0)
+	pair = append(pair, fastgpt.DataQuestion{
+		Q: *req.Q,
+		A: *req.A,
+	})
+	params.Data = append(params.Data, pair...)
+
+	resp, err := fastgpt.CreateBulkData(&params)
+	if err != nil {
+		return nil, errors.New("fastgpt create data failed")
+	}
+
+	if resp.Code == 0 {
+		return &types.BaseDataInfo{
+			Code: 0,
+			Msg:  errormsg.Success,
+			Data: fmt.Sprintf("Insert %d rows", resp.Data.InsertLen),
+		}, nil
+	}
+
+	return nil, errors.New(resp.StatusText)
+}

internal/logic/agent/create_agent_logic.go

@@ -2,6 +2,8 @@ package agent
 
 import (
 	"context"
+	"errors"
+	"wechat-api/hook/fastgpt"
 	"wechat-api/internal/svc"
 	"wechat-api/internal/types"
 	"wechat-api/internal/utils/dberrorhandler"
@@ -28,7 +30,7 @@ func NewCreateAgentLogic(ctx context.Context, svcCtx *svc.ServiceContext) *Creat
 func (l *CreateAgentLogic) CreateAgent(req *types.AgentInfo) (*types.BaseMsgResp, error) {
 	organizationId := l.ctx.Value("organizationId").(uint64)
 
-	_, err := l.svcCtx.DB.Agent.Create().
+	agent, err := l.svcCtx.DB.Agent.Create().
 		SetNotNilName(req.Name).
 		SetNotNilRole(req.Role).
 		//SetNotNilStatus(req.Status).
@@ -41,5 +43,44 @@ func (l *CreateAgentLogic) CreateAgent(req *types.AgentInfo) (*types.BaseMsgResp
 		return nil, dberrorhandler.DefaultEntError(l.Logger, err, req)
 	}
 
+	/**
+	1. 创建知识库
+	2. 创建默认集合
+	3. 保存到Agent信息里
+	*/
+	var datasetReq fastgpt.DatasetReq
+	datasetReq.Name = *req.Name
+	datasetReq.Intro = *req.Name
+	datasetReq.AgentModel = "gpt-3.5-turbo"
+	datasetReq.VectorModel = "text-embedding-ada-002"
+	datasetResp, err := fastgpt.CreateDataset(&datasetReq)
+	if err != nil {
+		return nil, errors.New("create dataset failed")
+	}
+	if datasetResp.Code != 0 {
+		return nil, errors.New(datasetResp.Message)
+	}
+
+	var collectionReq fastgpt.CreateCollectionReq
+	collectionReq.DatasetId = datasetResp.Data
+	collectionReq.Name = "手动录入"
+	collectionReq.Type = "folder"
+	collectionResp, err := fastgpt.CreateEmptyCollection(&collectionReq)
+	if err != nil {
+		return nil, errors.New("create dataset failed")
+	}
+	if collectionResp.Code != 0 {
+		return nil, errors.New(collectionResp.Message)
+	}
+
+	_, err = l.svcCtx.DB.Agent.UpdateOneID(agent.ID).
+		SetNotNilDatasetID(&datasetResp.Data).
+		SetNotNilCollectionID(&collectionResp.Data).
+		Save(l.ctx)
+
+	if err != nil {
+		return nil, errors.New("update dataset and collection failed")
+	}
+
 	return &types.BaseMsgResp{Msg: errormsg.CreateSuccess}, nil
 }

internal/logic/agent/get_agent_collection_list_logic.go

@@ -0,0 +1,70 @@
+package agent
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"wechat-api/hook/fastgpt"
+
+	"wechat-api/internal/svc"
+	"wechat-api/internal/types"
+
+	"github.com/zeromicro/go-zero/core/logx"
+)
+
+type GetAgentCollectionListLogic struct {
+	logx.Logger
+	ctx    context.Context
+	svcCtx *svc.ServiceContext
+}
+
+func NewGetAgentCollectionListLogic(ctx context.Context, svcCtx *svc.ServiceContext) *GetAgentCollectionListLogic {
+	return &GetAgentCollectionListLogic{
+		Logger: logx.WithContext(ctx),
+		ctx:    ctx,
+		svcCtx: svcCtx}
+}
+
+func (l *GetAgentCollectionListLogic) GetAgentCollectionList(req *types.CollectionListReq) (*types.CollectionListResp, error) {
+	collectionResp := types.CollectionListResp{}
+
+	var params fastgpt.GetCollectionListReq
+	params.DatasetId = *req.DatasetId
+	params.PageNum = *req.PageNum
+	params.PageSize = *req.PageNum
+	resp, err := fastgpt.GetCollectionList(&params)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, val := range resp.Data.Data {
+		fmt.Printf("id=%v, parentId=%v, total=%v \n", val.ID, val.ParentID, resp.Data.Total)
+	}
+
+	fmt.Printf("code=%v, message=%v, data=%v len=%v\n", resp.Code, resp.Message, resp.Data, len(resp.Data.Data))
+
+	if resp.Code == 200 && len(resp.Data.Data) > 0 {
+		collectionResp.Data = make([]types.CollectionSimpleInfo, 0, len(resp.Data.Data))
+		for _, val := range resp.Data.Data {
+			fmt.Printf("id=%v, parentId=%v \n", val.ID, val.ParentID)
+			collectionResp.Data = append(collectionResp.Data, types.CollectionSimpleInfo{
+				ID:             &val.ID,
+				ParentID:       &val.ParentID,
+				TmbId:          &val.TmbID,
+				Type:           &val.Type,
+				Name:           &val.Name,
+				DataAmount:     &val.DataAmount,
+				TrainingAmount: &val.TrainingAmount,
+			})
+		}
+
+		//collectionResp.Data = listData
+		collectionResp.PageNum = req.PageNum
+		collectionResp.PageSize = req.PageSize
+		collectionResp.Total = &resp.Data.Total
+
+		return &collectionResp, nil
+	}
+
+	return nil, errors.New(resp.Message)
+}

internal/logic/agent/get_agent_data_list_logic.go

@@ -0,0 +1,54 @@
+package agent
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"wechat-api/hook/fastgpt"
+
+	"wechat-api/internal/svc"
+	"wechat-api/internal/types"
+
+	"github.com/zeromicro/go-zero/core/logx"
+)
+
+type GetAgentDataListLogic struct {
+	logx.Logger
+	ctx    context.Context
+	svcCtx *svc.ServiceContext
+}
+
+func NewGetAgentDataListLogic(ctx context.Context, svcCtx *svc.ServiceContext) *GetAgentDataListLogic {
+	return &GetAgentDataListLogic{
+		Logger: logx.WithContext(ctx),
+		ctx:    ctx,
+		svcCtx: svcCtx}
+}
+
+func (l *GetAgentDataListLogic) GetAgentDataList(req *types.DataListReq) (*types.DataListResp, error) {
+	var params fastgpt.GetDataListReq
+	params.CollectionId = *req.CollectionId
+	params.PageSize = *req.PageSize
+	params.PageNum = *req.PageNum
+	resp, err := fastgpt.GetDataList(&params)
+	if err != nil {
+		fmt.Printf("%v", err.Error())
+	}
+	dataResp := types.DataListResp{}
+	if resp.Code == 0 {
+		for _, val := range resp.Data.Data {
+			dataResp.Data = append(dataResp.Data, types.DataInfo{
+				ID:           &val.ID,
+				Q:            &val.Q,
+				A:            &val.A,
+				ChunkIndex:   &val.ChunkIndex,
+				DatasetId:    &val.DatasetID,
+				CollectionId: &val.CollectionID,
+			})
+		}
+
+		return &dataResp, nil
+	}
+
+	return nil, errors.New(resp.Message)
+}

internal/logic/agent/update_agent_data_logic.go

@@ -0,0 +1,47 @@
+package agent
+
+import (
+	"context"
+	"errors"
+	"github.com/suyuan32/simple-admin-common/msg/errormsg"
+	"wechat-api/hook/fastgpt"
+
+	"wechat-api/internal/svc"
+	"wechat-api/internal/types"
+
+	"github.com/zeromicro/go-zero/core/logx"
+)
+
+type UpdateAgentDataLogic struct {
+	logx.Logger
+	ctx    context.Context
+	svcCtx *svc.ServiceContext
+}
+
+func NewUpdateAgentDataLogic(ctx context.Context, svcCtx *svc.ServiceContext) *UpdateAgentDataLogic {
+	return &UpdateAgentDataLogic{
+		Logger: logx.WithContext(ctx),
+		ctx:    ctx,
+		svcCtx: svcCtx}
+}
+
+func (l *UpdateAgentDataLogic) UpdateAgentData(req *types.UpdateDataInfoReq) (*types.BaseDataInfo, error) {
+	params := fastgpt.UpdateDataReq{}
+	params.ID = *req.ID
+	params.Q = *req.Q
+	params.A = *req.A
+	resp, err := fastgpt.UpdateData(&params)
+	if err != nil {
+		return nil, errors.New("fastgpt update data failed")
+	}
+
+	if resp.Code == 0 {
+		return &types.BaseDataInfo{
+			Code: 0,
+			Msg:  errormsg.Success,
+			Data: "",
+		}, nil
+	}
+
+	return nil, errors.New(resp.StatusText)
+}

internal/types/types.go

@@ -423,6 +423,194 @@ type AgentInfoResp struct {
 	Data AgentInfo `json:"data"`
 }
 
+type VectorModel struct {
+	Model            *string `json:"model"`
+	Name             *string `json:"name"`
+	CharsPointsPrice *uint64 `json:"charsPointsPrice"`
+	DefaultToken     *uint64 `json:"defaultToken"`
+	MaxToken         *uint64 `json:"maxToken"`
+	Weight           *uint64 `json:"weight"`
+}
+
+type AgentModel struct {
+	Model            *string `json:"model"`
+	Name             *string `json:"name"`
+	MaxContext       *uint64 `json:"maxContext"`
+	MaxResponse      *uint64 `json:"maxResponse"`
+	CharsPointsPrice *uint64 `json:"charsPointsPrice"`
+}
+
+type DatasetId struct {
+	ID          *string `json:"id"`
+	ParentID    *string `json:"parentId"`
+	TeamId      *string `json:"teamId"`
+	TmbId       *string `json:"tmbId"`
+	Type        *string `json:"type"`
+	Status      *string `json:"status"`
+	Name        *string `json:"name"`
+	VectorModel *string `json:"vectorModel"`
+	AgentModel  *string `json:"agentModel"`
+	Intro       *string `json:"intro"`
+	Permission  *string `json:"permission"`
+}
+
+// Dataset info | 知识库详情
+// swagger:model DatasetInfo
+type DatasetInfo struct {
+	ID          *string     `json:"id"`
+	ParentID    *string     `json:"parentId"`
+	TeamId      *string     `json:"teamId"`
+	TmbId       *string     `json:"tmbId"`
+	Type        *string     `json:"type"`
+	Name        *string     `json:"name"`
+	Intro       *string     `json:"intro"`
+	Status      *string     `json:"status"`
+	Avatar      *string     `json:"avatar"`
+	VectorModel VectorModel `json:"vectorModel"`
+	AgentModel  AgentModel  `json:"agentModel"`
+	Permission  *string     `json:"permission"`
+	CanWrite    *bool       `json:"canWrite"`
+	IsOwner     *bool       `json:"isOwner"`
+}
+
+type Index struct {
+	DefaultIndex *bool   `json:"canWrite"`
+	Type         *string `json:"type"`
+	DataId       *string `json:"dataId"`
+	Text         *string `json:"text"`
+	ID           *string `json:"id"`
+}
+
+// Collection Info | 集合详情
+// swagger:model CollectionInfo
+type CollectionInfo struct {
+	ID             *string   `json:"id"`
+	ParentID       *string   `json:"parentId"`
+	TmbId          *string   `json:"tmbId,optional"`
+	Type           *string   `json:"type"`
+	Name           *string   `json:"name"`
+	DataAmount     *uint64   `json:"dataAmount,optional"`
+	TrainingAmount *uint64   `json:"trainingAmount,optional"`
+	TrainingType   *string   `json:"trainingType,optional"`
+	ChunkSize      *uint64   `json:"chunkSize,optional"`
+	ChunkSplitter  *string   `json:"chunkSplitter,optional"`
+	QaPrompt       *string   `json:"qaPrompt,optional"`
+	RawTextLength  *uint64   `json:"rawTextLength,optional"`
+	CanWrite       *bool     `json:"canWrite,optional"`
+	SourceName     *string   `json:"sourceName,optional"`
+	DatasetId      DatasetId `json:"datasetId,optional"`
+}
+
+// swagger:model CollectionSimpleInfo
+type CollectionSimpleInfo struct {
+	ID             *string `json:"id"`
+	ParentID       *string `json:"parentId"`
+	TmbId          *string `json:"tmbId,optional"`
+	Type           *string `json:"type"`
+	Name           *string `json:"name"`
+	DataAmount     *uint64 `json:"dataAmount,optional"`
+	TrainingAmount *uint64 `json:"trainingAmount,optional"`
+}
+
+// swagger:model DataInfo
+type DataInfo struct {
+	ID           *string `json:"id"`
+	Q            *string `json:"q"`
+	A            *string `json:"a"`
+	ChunkIndex   *uint64 `json:"chunkIndex"`
+	Indexes      []Index `json:"indexes"`
+	DatasetId    *string `json:"datasetId"`
+	CollectionId *string `json:"collectionId"`
+	SourceName   *string `json:"sourceName"`
+	SourceId     *string `json:"sourceId"`
+	CanWrite     *bool   `json:"canWrite"`
+	IsOwner      *bool   `json:"isOwner"`
+}
+
+// Get collection list request params | Collection列表请求参数
+// swagger:model CollectionListReq
+type CollectionListReq struct {
+	// required : true
+	// min : 0
+	PageNum *int `json:"pageNum" validate:"required,number,gt=0"`
+	// required : true
+	// max : 100000
+	PageSize *int `json:"pageSize" validate:"required,number,lt=100000"`
+	// required : true
+	DatasetId *string `json:"datasetId" validate:"required"`
+}
+
+// Collection list response | Collection List信息返回体
+// swagger:model CollectionListResp
+type CollectionListResp struct {
+	BaseDataInfo
+	// Agent information | Agent数据
+	Data []CollectionSimpleInfo `json:"data"`
+	// required : true
+	// min : 0
+	PageNum *int `json:"pageNum" validate:"required,number,gt=0"`
+	// required : true
+	// max : 100000
+	PageSize *int `json:"pageSize" validate:"required,number,lt=100000"`
+	Total    *int `json:"total"`
+}
+
+// Get collection list request params | Collection列表请求参数
+// swagger:model DataListReq
+type DataListReq struct {
+	// required : true
+	// min : 0
+	PageNum *int `json:"pageNum" validate:"required,number,gt=0"`
+	// required : true
+	// max : 100000
+	PageSize *int `json:"pageSize" validate:"required,number,lt=100000"`
+	// required : true
+	CollectionId *string `json:"collectionId" validate:"required"`
+}
+
+// Data list response | Data List信息返回体
+// swagger:model DataListResp
+type DataListResp struct {
+	BaseDataInfo
+	// Agent information | Agent数据
+	Data []DataInfo `json:"data"`
+	// required : true
+	// min : 0
+	PageNum *int `json:"pageNum" validate:"required,number,gt=0"`
+	// required : true
+	// max : 100000
+	PageSize *int `json:"pageSize" validate:"required,number,lt=100000"`
+	Total    *int `json:"total"`
+}
+
+// Data create request | 信息请求体
+// swagger:model CreateDataInfoReq
+type CreateDataInfoReq struct {
+	// required : true
+	CollectionId *string `json:"collectionId" validate:"required"`
+	// Q
+	Q *string `json:"q"`
+	// A
+	A *string `json:"a"`
+}
+
+//	IndexSingle {
+//		Text A *string `json:"text"`
+//	}
+//
+// Data create request | 信息返回体
+// swagger:model UpdateDataInfoReq
+type UpdateDataInfoReq struct {
+	// required : true
+	CollectionId *string `json:"collectionId" validate:"required"`
+	// ID
+	ID *string `json:"id"`
+	// Q
+	Q *string `json:"q"`
+	// A
+	A *string `json:"a"`
+}
+
 type LoginQRStatus struct {
 	// 登陆二维码
 	QRCode string `json:"qRCode,optional"`