estel/dify
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

feat: mypy for all type check (#10921)

Browse Source
This commit is contained in:
yihong
2024-12-24 18:38:51 +08:00
committed by GitHub
parent c91e8b1737
commit 56e15d09a9
584 changed files with 3975 additions and 2826 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
api/models/dataset.py
View File

@@ -9,6 +9,7 @@ import pickle
import re
import time
from json import JSONDecodeError
from typing import Any, cast
from sqlalchemy import func
from sqlalchemy.dialects.postgresql import JSONB
@@ -29,7 +30,7 @@ class DatasetPermissionEnum(enum.StrEnum):
PARTIAL_TEAM = "partial_members"
class Dataset(db.Model):
class Dataset(db.Model): # type: ignore[name-defined]
__tablename__ = "datasets"
__table_args__ = (
db.PrimaryKeyConstraint("id", name="dataset_pkey"),
@@ -200,7 +201,7 @@ class Dataset(db.Model):
return f"Vector_index_{normalized_dataset_id}_Node"
class DatasetProcessRule(db.Model):
class DatasetProcessRule(db.Model): # type: ignore[name-defined]
__tablename__ = "dataset_process_rules"
__table_args__ = (
db.PrimaryKeyConstraint("id", name="dataset_process_rule_pkey"),
@@ -216,7 +217,7 @@ class DatasetProcessRule(db.Model):
MODES = ["automatic", "custom"]
PRE_PROCESSING_RULES = ["remove_stopwords", "remove_extra_spaces", "remove_urls_emails"]
AUTOMATIC_RULES = {
AUTOMATIC_RULES: dict[str, Any] = {
"pre_processing_rules": [
{"id": "remove_extra_spaces", "enabled": True},
{"id": "remove_urls_emails", "enabled": False},
@@ -242,7 +243,7 @@ class DatasetProcessRule(db.Model):
return None
class Document(db.Model):
class Document(db.Model): # type: ignore[name-defined]
__tablename__ = "documents"
__table_args__ = (
db.PrimaryKeyConstraint("id", name="document_pkey"),
@@ -492,7 +493,7 @@ class Document(db.Model):
)
class DocumentSegment(db.Model):
class DocumentSegment(db.Model): # type: ignore[name-defined]
__tablename__ = "document_segments"
__table_args__ = (
db.PrimaryKeyConstraint("id", name="document_segment_pkey"),
@@ -604,7 +605,7 @@ class DocumentSegment(db.Model):
return text
class AppDatasetJoin(db.Model):
class AppDatasetJoin(db.Model): # type: ignore[name-defined]
__tablename__ = "app_dataset_joins"
__table_args__ = (
db.PrimaryKeyConstraint("id", name="app_dataset_join_pkey"),
@@ -621,7 +622,7 @@ class AppDatasetJoin(db.Model):
return db.session.get(App, self.app_id)
class DatasetQuery(db.Model):
class DatasetQuery(db.Model): # type: ignore[name-defined]
__tablename__ = "dataset_queries"
__table_args__ = (
db.PrimaryKeyConstraint("id", name="dataset_query_pkey"),
@@ -638,7 +639,7 @@ class DatasetQuery(db.Model):
created_at = db.Column(db.DateTime, nullable=False, server_default=db.func.current_timestamp())
class DatasetKeywordTable(db.Model):
class DatasetKeywordTable(db.Model): # type: ignore[name-defined]
__tablename__ = "dataset_keyword_tables"
__table_args__ = (
db.PrimaryKeyConstraint("id", name="dataset_keyword_table_pkey"),
@@ -683,7 +684,7 @@ class DatasetKeywordTable(db.Model):
return None
class Embedding(db.Model):
class Embedding(db.Model): # type: ignore[name-defined]
__tablename__ = "embeddings"
__table_args__ = (
db.PrimaryKeyConstraint("id", name="embedding_pkey"),
@@ -704,10 +705,10 @@ class Embedding(db.Model):
self.embedding = pickle.dumps(embedding_data, protocol=pickle.HIGHEST_PROTOCOL)
def get_embedding(self) -> list[float]:
return pickle.loads(self.embedding)
return cast(list[float], pickle.loads(self.embedding))
class DatasetCollectionBinding(db.Model):
class DatasetCollectionBinding(db.Model): # type: ignore[name-defined]
__tablename__ = "dataset_collection_bindings"
__table_args__ = (
db.PrimaryKeyConstraint("id", name="dataset_collection_bindings_pkey"),
@@ -722,7 +723,7 @@ class DatasetCollectionBinding(db.Model):
created_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())
class TidbAuthBinding(db.Model):
class TidbAuthBinding(db.Model): # type: ignore[name-defined]
__tablename__ = "tidb_auth_bindings"
__table_args__ = (
db.PrimaryKeyConstraint("id", name="tidb_auth_bindings_pkey"),
@@ -742,7 +743,7 @@ class TidbAuthBinding(db.Model):
created_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())
class Whitelist(db.Model):
class Whitelist(db.Model): # type: ignore[name-defined]
__tablename__ = "whitelists"
__table_args__ = (
db.PrimaryKeyConstraint("id", name="whitelists_pkey"),
@@ -754,7 +755,7 @@ class Whitelist(db.Model):
created_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())
class DatasetPermission(db.Model):
class DatasetPermission(db.Model): # type: ignore[name-defined]
__tablename__ = "dataset_permissions"
__table_args__ = (
db.PrimaryKeyConstraint("id", name="dataset_permission_pkey"),
@@ -771,7 +772,7 @@ class DatasetPermission(db.Model):
created_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())
class ExternalKnowledgeApis(db.Model):
class ExternalKnowledgeApis(db.Model): # type: ignore[name-defined]
__tablename__ = "external_knowledge_apis"
__table_args__ = (
db.PrimaryKeyConstraint("id", name="external_knowledge_apis_pkey"),
@@ -824,7 +825,7 @@ class ExternalKnowledgeApis(db.Model):
return dataset_bindings
class ExternalKnowledgeBindings(db.Model):
class ExternalKnowledgeBindings(db.Model): # type: ignore[name-defined]
__tablename__ = "external_knowledge_bindings"
__table_args__ = (
db.PrimaryKeyConstraint("id", name="external_knowledge_bindings_pkey"),