67 lines
1.7 KiB
Python
67 lines
1.7 KiB
Python
from pydantic import BaseModel, Field
|
||
from typing import Literal
|
||
|
||
|
||
class TagProperties(BaseModel):
|
||
tag_path: str = Field(..., description="完整标签路径，如 '编程语言/Python/异步'")
|
||
short_name: str = Field(..., description="显示名称，如 '异步'")
|
||
level: int = Field(..., ge=1, description="层级深度，1为顶级")
|
||
parent_path: str | None = Field(None, description="父路径，如 '编程语言/Python'")
|
||
description: str | None = Field(None, description="AI生成的标签描述")
|
||
color: str | None = Field(None, description="标签颜色，如 '#FF5733'")
|
||
|
||
|
||
class KGNodeOut(BaseModel):
|
||
id: str
|
||
name: str
|
||
entity_type: str
|
||
description: str | None
|
||
properties_: dict | None
|
||
importance: float
|
||
created_at: str
|
||
# 新增：如果是 tag 节点，返回 tag 属性
|
||
tag_properties: TagProperties | None = None
|
||
|
||
model_config = {"from_attributes": True}
|
||
|
||
def model_post_init(self, __context):
|
||
if self.entity_type == "tag" and self.properties_:
|
||
self.tag_properties = TagProperties(**self.properties_)
|
||
|
||
|
||
class KGEdgeOut(BaseModel):
|
||
id: str
|
||
source_id: str
|
||
target_id: str
|
||
relation_type: str
|
||
weight: float
|
||
properties_: dict | None
|
||
|
||
model_config = {"from_attributes": True}
|
||
|
||
|
||
class GraphOut(BaseModel):
|
||
nodes: list[KGNodeOut]
|
||
edges: list[KGEdgeOut]
|
||
|
||
|
||
class KGBuildRequest(BaseModel):
|
||
user_id: str
|
||
document_ids: list[str] | None = None # None = 全量重建
|
||
|
||
|
||
class TagExtractRequest(BaseModel):
|
||
content: str = Field(..., min_length=10)
|
||
user_id: str
|
||
|
||
|
||
class TagExtractResponse(BaseModel):
|
||
tags: list[TagProperties]
|
||
tag_count: int
|
||
|
||
|
||
class RelatedContentRequest(BaseModel):
|
||
tag_ids: list[str]
|
||
user_id: str
|
||
limit: int = 10
|