diff --git a/.dockerignore b/.dockerignore index bae9066..5ad4e6e 100644 --- a/.dockerignore +++ b/.dockerignore @@ -15,7 +15,9 @@ __pycache__ htmlcov data logs -models +# Root weights directory only. A bare name would also drop Python packages +# such as kamidso/engine/serve/rag/models. +/models plugins pilot/data pilot/message diff --git a/.gitignore b/.gitignore index 2d1af8e..5703e96 100644 --- a/.gitignore +++ b/.gitignore @@ -37,7 +37,9 @@ pilot/meta_data/*.db pilot/meta_data/*.db-* pilot/meta_data/alembic/versions/* pilot/benchmark_meta_data/*.db -models/ +# Root weights directory only. A bare `models/` also drops Python packages +# such as kamidso/engine/serve/rag/models. +/models/ nltk_data/ .chroma/ .vectordb/ diff --git a/kamidso/engine/serve/agent/chat/models/__init__.py b/kamidso/engine/serve/agent/chat/models/__init__.py new file mode 100644 index 0000000..f5ca4a0 --- /dev/null +++ b/kamidso/engine/serve/agent/chat/models/__init__.py @@ -0,0 +1,2 @@ +# This is an auto-generated __init__.py file +# generated by `kamidso.engine.core new serve agent/chat` diff --git a/kamidso/engine/serve/agent/chat/models/models.py b/kamidso/engine/serve/agent/chat/models/models.py new file mode 100644 index 0000000..d2aa721 --- /dev/null +++ b/kamidso/engine/serve/agent/chat/models/models.py @@ -0,0 +1,75 @@ +"""This is an auto-generated model file +You can define your own models and DAOs here +""" + +from datetime import datetime +from typing import Any, Dict, Union + +from sqlalchemy import Column, DateTime, Integer + +from kamidso.engine.core.storage.metadata import BaseDao, Model + +from ..api.schemas import ServeRequest, ServerResponse +from ..config import SERVER_APP_TABLE_NAME, ServeConfig + + +class ServeEntity(Model): + __tablename__ = SERVER_APP_TABLE_NAME + id = Column(Integer, primary_key=True, comment="Auto increment id") + + # TODO: define your own fields here + + gmt_created = Column(DateTime, default=datetime.now, comment="Record creation time") + gmt_modified = Column(DateTime, default=datetime.now, comment="Record update time") + + def __repr__(self): + return ( + f"ServeEntity(id={self.id}, gmt_created='{self.gmt_created}', " + f"gmt_modified='{self.gmt_modified}')" + ) + + +class ServeDao(BaseDao[ServeEntity, ServeRequest, ServerResponse]): + """The DAO class for Agent/chat""" + + def __init__(self, serve_config: ServeConfig): + super().__init__() + self._serve_config = serve_config + + def from_request(self, request: Union[ServeRequest, Dict[str, Any]]) -> ServeEntity: + """Convert the request to an entity + + Args: + request (Union[ServeRequest, Dict[str, Any]]): The request + + Returns: + T: The entity + """ + request_dict = request.dict() if isinstance(request, ServeRequest) else request + entity = ServeEntity(**request_dict) + # TODO implement your own logic here, transfer the request_dict to an entity + return entity + + def to_request(self, entity: ServeEntity) -> ServeRequest: + """Convert the entity to a request + + Args: + entity (T): The entity + + Returns: + REQ: The request + """ + # TODO implement your own logic here, transfer the entity to a request + return ServeRequest() + + def to_response(self, entity: ServeEntity) -> ServerResponse: + """Convert the entity to a response + + Args: + entity (T): The entity + + Returns: + RES: The response + """ + # TODO implement your own logic here, transfer the entity to a response + return ServerResponse() diff --git a/kamidso/engine/serve/connector/models/__init__.py b/kamidso/engine/serve/connector/models/__init__.py new file mode 100644 index 0000000..b767de5 --- /dev/null +++ b/kamidso/engine/serve/connector/models/__init__.py @@ -0,0 +1 @@ +"""Connector models module.""" diff --git a/kamidso/engine/serve/connector/models/models.py b/kamidso/engine/serve/connector/models/models.py new file mode 100644 index 0000000..748ba6c --- /dev/null +++ b/kamidso/engine/serve/connector/models/models.py @@ -0,0 +1,135 @@ +import logging +import uuid +from datetime import datetime +from typing import Any, Dict, Union + +from sqlalchemy import Column, DateTime, Integer, String, Text + +from kamidso.engine.core.storage.metadata import BaseDao, Model + +from ..config import ServeConfig + +logger = logging.getLogger(__name__) + + +class ConnectorInstanceEntity(Model): + __tablename__ = "connector_instance" + + id = Column( + Integer, primary_key=True, autoincrement=True, comment="Auto increment id" + ) + connector_id = Column( + String(64), unique=True, nullable=False, comment="Connector UUID" + ) + connector_type = Column( + String(64), nullable=False, comment="Connector type, e.g. yuque, feishu" + ) + display_name = Column(String(256), nullable=True, comment="Display name") + encrypted_credentials = Column( + Text, nullable=True, comment="Encrypted credentials JSON" + ) + encryption_salt = Column(String(256), nullable=True, comment="Encryption salt") + status = Column( + String(32), + nullable=True, + comment="Status: active/error/disconnected/needs_reactivation", + ) + config_json = Column(Text, nullable=True, comment="Optional extra config JSON") + user_name = Column(String(128), index=True, nullable=True, comment="User name") + sys_code = Column(String(128), index=True, nullable=True, comment="System code") + gmt_created = Column(DateTime, default=datetime.now, comment="Record creation time") + gmt_modified = Column( + DateTime, + default=datetime.now, + onupdate=datetime.now, + comment="Record update time", + ) + + def __repr__(self): + return ( + f"ConnectorInstanceEntity(id={self.id}, " + f"connector_id='{self.connector_id}', " + f"connector_type='{self.connector_type}', " + f"display_name='{self.display_name}', " + f"status='{self.status}', user_name='{self.user_name}')" + ) + + +class ConnectorInstanceDao( + BaseDao[ConnectorInstanceEntity, Dict[str, Any], Dict[str, Any]] +): + """DAO for ConnectorInstanceEntity providing basic CRUD operations.""" + + def __init__(self, serve_config: ServeConfig): + super().__init__() + self._serve_config = serve_config + + def from_request( + self, request: Union[Dict[str, Any], Any] + ) -> ConnectorInstanceEntity: + """Convert a request dict to a ConnectorInstanceEntity. + + Args: + request: Request dict or object with connector fields. + + Returns: + ConnectorInstanceEntity: The entity instance. + """ + request_dict = request if isinstance(request, dict) else request.dict() + entity = ConnectorInstanceEntity(**request_dict) + if not entity.connector_id: + entity.connector_id = str(uuid.uuid4()) + return entity + + def to_request(self, entity: ConnectorInstanceEntity) -> Dict[str, Any]: + """Convert a ConnectorInstanceEntity to a request dict. + + Args: + entity: The entity instance. + + Returns: + Dict[str, Any]: The request dict. + """ + return { + "connector_id": entity.connector_id, + "connector_type": entity.connector_type, + "display_name": entity.display_name, + "encrypted_credentials": entity.encrypted_credentials, + "encryption_salt": entity.encryption_salt, + "status": entity.status, + "config_json": entity.config_json, + "user_name": entity.user_name, + "sys_code": entity.sys_code, + } + + def to_response(self, entity: ConnectorInstanceEntity) -> Dict[str, Any]: + """Convert a ConnectorInstanceEntity to a response dict. + + Args: + entity: The entity instance. + + Returns: + Dict[str, Any]: The response dict. + """ + gmt_created_str = ( + entity.gmt_created.strftime("%Y-%m-%d %H:%M:%S") + if entity.gmt_created + else None + ) + gmt_modified_str = ( + entity.gmt_modified.strftime("%Y-%m-%d %H:%M:%S") + if entity.gmt_modified + else None + ) + return { + "id": entity.id, + "connector_id": entity.connector_id, + "connector_type": entity.connector_type, + "display_name": entity.display_name, + "status": entity.status, + "config_json": entity.config_json, + "user_name": entity.user_name, + "sys_code": entity.sys_code, + "gmt_created": gmt_created_str, + "gmt_modified": gmt_modified_str, + } diff --git a/kamidso/engine/serve/conversation/models/__init__.py b/kamidso/engine/serve/conversation/models/__init__.py new file mode 100644 index 0000000..0d4da22 --- /dev/null +++ b/kamidso/engine/serve/conversation/models/__init__.py @@ -0,0 +1,2 @@ +# This is an auto-generated __init__.py file +# generated by `kamidso.engine.core new serve conversation` diff --git a/kamidso/engine/serve/conversation/models/models.py b/kamidso/engine/serve/conversation/models/models.py new file mode 100644 index 0000000..4117601 --- /dev/null +++ b/kamidso/engine/serve/conversation/models/models.py @@ -0,0 +1,168 @@ +"""This is an auto-generated model file +You can define your own models and DAOs here +""" + +import json +from typing import Any, Dict, List, Optional, Union + +from kamidso.engine.core.core import MessageStorageItem +from kamidso.engine.core.storage.chat_history.chat_history_db import ChatHistoryEntity as ServeEntity +from kamidso.engine.core.storage.chat_history.chat_history_db import ChatHistoryMessageEntity +from kamidso.engine.core.storage.metadata import BaseDao +from kamidso.engine.core.util import PaginationResult + +from ..api.schemas import ServeRequest, ServerResponse +from ..config import ServeConfig + + +class ServeDao(BaseDao[ServeEntity, ServeRequest, ServerResponse]): + """The DAO class for Conversation""" + + def __init__(self, serve_config: ServeConfig): + super().__init__() + self._serve_config = serve_config + + def from_request(self, request: Union[ServeRequest, Dict[str, Any]]) -> ServeEntity: + """Convert the request to an entity + + Args: + request (Union[ServeRequest, Dict[str, Any]]): The request + + Returns: + T: The entity + """ + request_dict = ( + request.to_dict() if isinstance(request, ServeRequest) else request + ) + entity = ServeEntity(**request_dict) + # TODO implement your own logic here, transfer the request_dict to an entity + return entity + + def to_request(self, entity: ServeEntity) -> ServeRequest: + """Convert the entity to a request + + Args: + entity (T): The entity + + Returns: + REQ: The request + """ + # TODO implement your own logic here, transfer the entity to a request + return ServeRequest() + + def to_response(self, entity: ServeEntity) -> ServerResponse: + """Convert the entity to a response + + Args: + entity (T): The entity + + Returns: + RES: The response + """ + # TODO implement your own logic here, transfer the entity to a response + gmt_created = ( + entity.gmt_created.strftime("%Y-%m-%d %H:%M:%S") + if entity.gmt_created + else None + ) + gmt_modified = ( + entity.gmt_modified.strftime("%Y-%m-%d %H:%M:%S") + if entity.gmt_modified + else None + ) + + return ServerResponse( + app_code=entity.app_code, + conv_uid=entity.conv_uid, + user_input=entity.summary, + chat_mode=entity.chat_mode, + user_name="", + sys_code=entity.sys_code, + gmt_created=gmt_created, + gmt_modified=gmt_modified, + ) + + def get_latest_message(self, conv_uid: str) -> Optional[MessageStorageItem]: + """Get the latest message of a conversation + + Args: + conv_uid (str): The conversation UID + + Returns: + ChatHistoryMessageEntity: The latest message + """ + with self.session() as session: + entity: ChatHistoryMessageEntity = ( + session.query(ChatHistoryMessageEntity) + .filter(ChatHistoryMessageEntity.conv_uid == conv_uid) + .order_by(ChatHistoryMessageEntity.gmt_created.desc()) + .first() + ) + if not entity: + return None + message_detail = ( + json.loads(entity.message_detail) if entity.message_detail else {} + ) + return MessageStorageItem(entity.conv_uid, entity.index, message_detail) + + def _parse_old_messages(self, entity: ServeEntity) -> List[Dict[str, Any]]: + """Parse the old messages + + Args: + entity (ServeEntity): The entity + + Returns: + str: The old messages + """ + messages = json.loads(entity.messages) + return messages + + def get_conv_by_page( + self, req: ServeRequest, page: int, page_size: int + ) -> PaginationResult[ServerResponse]: + """Get conversation by page + + Args: + req (ServeRequest): The request + page (int): The page number + page_size (int): The page size + + Returns: + List[ChatHistoryEntity]: The conversation list + """ + with self.session(commit=False) as session: + query = self._create_query_object(session, req) + query = query.order_by(ServeEntity.gmt_created.desc()) + total_count = query.count() + items = query.offset((page - 1) * page_size).limit(page_size) + total_pages = (total_count + page_size - 1) // page_size + result_items = [] + for item in items: + select_param, model_name = "", None + if item.messages: + messages = self._parse_old_messages(item) + last_round = max(messages, key=lambda x: x["chat_order"]) + if "param_value" in last_round: + select_param = last_round["param_value"] + else: + select_param = "" + else: + latest_message = self.get_latest_message(item.conv_uid) + if latest_message: + message = latest_message.to_message() + select_param = message.additional_kwargs.get("param_value") + model_name = message.additional_kwargs.get("model_name") + res_item = self.to_response(item) + res_item.select_param = select_param + res_item.model_name = model_name + result_items.append(res_item) + + result = PaginationResult( + items=result_items, + total_count=total_count, + total_pages=total_pages, + page=page, + page_size=page_size, + ) + + return result diff --git a/kamidso/engine/serve/datasource/models/__init__.py b/kamidso/engine/serve/datasource/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/kamidso/engine/serve/datasource/models/models.py b/kamidso/engine/serve/datasource/models/models.py new file mode 100644 index 0000000..e69de29 diff --git a/kamidso/engine/serve/file/models/__init__.py b/kamidso/engine/serve/file/models/__init__.py new file mode 100644 index 0000000..1341010 --- /dev/null +++ b/kamidso/engine/serve/file/models/__init__.py @@ -0,0 +1,2 @@ +# This is an auto-generated __init__.py file +# generated by `kamidso.engine.core new serve file` diff --git a/kamidso/engine/serve/file/models/file_adapter.py b/kamidso/engine/serve/file/models/file_adapter.py new file mode 100644 index 0000000..c7a28c1 --- /dev/null +++ b/kamidso/engine/serve/file/models/file_adapter.py @@ -0,0 +1,88 @@ +import json +from typing import Type + +from sqlalchemy.orm import Session + +from kamidso.engine.core.core.interface.file import FileMetadata, FileMetadataIdentifier +from kamidso.engine.core.core.interface.storage import StorageItemAdapter + +from .models import ServeEntity + + +class FileMetadataAdapter(StorageItemAdapter[FileMetadata, ServeEntity]): + """File metadata adapter. + + Convert between storage format and database model. + """ + + def to_storage_format(self, item: FileMetadata) -> ServeEntity: + """Convert to storage format.""" + custom_metadata = ( + {k: v for k, v in item.custom_metadata.items()} + if item.custom_metadata + else {} + ) + user_name = item.user_name or custom_metadata.get("user_name") + sys_code = item.sys_code or custom_metadata.get("sys_code") + if "user_name" in custom_metadata: + del custom_metadata["user_name"] + if "sys_code" in custom_metadata: + del custom_metadata["sys_code"] + custom_metadata_json = ( + json.dumps(custom_metadata, ensure_ascii=False) if custom_metadata else None + ) + return ServeEntity( + bucket=item.bucket, + file_id=item.file_id, + file_name=item.file_name, + file_size=item.file_size, + storage_type=item.storage_type, + storage_path=item.storage_path, + uri=item.uri, + custom_metadata=custom_metadata_json, + file_hash=item.file_hash, + user_name=user_name, + sys_code=sys_code, + ) + + def from_storage_format(self, model: ServeEntity) -> FileMetadata: + """Convert from storage format.""" + custom_metadata = ( + json.loads(model.custom_metadata) if model.custom_metadata else None + ) + if custom_metadata is None: + custom_metadata = {} + if model.user_name: + custom_metadata["user_name"] = model.user_name + if model.sys_code: + custom_metadata["sys_code"] = model.sys_code + + return FileMetadata( + bucket=model.bucket, + file_id=model.file_id, + file_name=model.file_name, + file_size=model.file_size, + storage_type=model.storage_type, + storage_path=model.storage_path, + uri=model.uri, + custom_metadata=custom_metadata, + file_hash=model.file_hash, + user_name=model.user_name, + sys_code=model.sys_code, + ) + + def get_query_for_identifier( + self, + storage_format: Type[ServeEntity], + resource_id: FileMetadataIdentifier, + **kwargs, + ): + """Get query for identifier.""" + session: Session = kwargs.get("session") + if session is None: + raise Exception("session is None") + return ( + session.query(storage_format) + .filter(storage_format.bucket == resource_id.bucket) + .filter(storage_format.file_id == resource_id.file_id) + ) diff --git a/kamidso/engine/serve/file/models/models.py b/kamidso/engine/serve/file/models/models.py new file mode 100644 index 0000000..8b42cca --- /dev/null +++ b/kamidso/engine/serve/file/models/models.py @@ -0,0 +1,90 @@ +"""This is an auto-generated model file +You can define your own models and DAOs here +""" + +from datetime import datetime +from typing import Any, Dict, Union + +from sqlalchemy import Column, DateTime, Integer, String, Text, UniqueConstraint + +from kamidso.engine.core.storage.metadata import BaseDao, Model + +from ..api.schemas import ServeRequest, ServerResponse +from ..config import SERVER_APP_TABLE_NAME, ServeConfig + + +class ServeEntity(Model): + __tablename__ = SERVER_APP_TABLE_NAME + __table_args__ = (UniqueConstraint("bucket", "file_id", name="uk_bucket_file_id"),) + + id = Column(Integer, primary_key=True, comment="Auto increment id") + + bucket = Column(String(255), nullable=False, comment="Bucket name") + file_id = Column(String(255), nullable=False, comment="File id") + file_name = Column(String(256), nullable=False, comment="File name") + file_size = Column(Integer, nullable=True, comment="File size") + storage_type = Column(String(32), nullable=False, comment="Storage type") + storage_path = Column(String(512), nullable=False, comment="Storage path") + uri = Column(String(512), nullable=False, comment="File URI") + custom_metadata = Column( + Text, nullable=True, comment="Custom metadata, JSON format" + ) + file_hash = Column(String(128), nullable=True, comment="File hash") + user_name = Column(String(128), index=True, nullable=True, comment="User name") + sys_code = Column(String(128), index=True, nullable=True, comment="System code") + gmt_created = Column(DateTime, default=datetime.now, comment="Record creation time") + gmt_modified = Column(DateTime, default=datetime.now, comment="Record update time") + + def __repr__(self): + return ( + f"ServeEntity(id={self.id}, gmt_created='{self.gmt_created}', " + f"gmt_modified='{self.gmt_modified}')" + ) + + +class ServeDao(BaseDao[ServeEntity, ServeRequest, ServerResponse]): + """The DAO class for File""" + + def __init__(self, serve_config: ServeConfig): + super().__init__() + self._serve_config = serve_config + + def from_request(self, request: Union[ServeRequest, Dict[str, Any]]) -> ServeEntity: + """Convert the request to an entity + + Args: + request (Union[ServeRequest, Dict[str, Any]]): The request + + Returns: + T: The entity + """ + request_dict = ( + request.to_dict() if isinstance(request, ServeRequest) else request + ) + entity = ServeEntity(**request_dict) + # TODO implement your own logic here, transfer the request_dict to an entity + return entity + + def to_request(self, entity: ServeEntity) -> ServeRequest: + """Convert the entity to a request + + Args: + entity (T): The entity + + Returns: + REQ: The request + """ + # TODO implement your own logic here, transfer the entity to a request + return ServeRequest() + + def to_response(self, entity: ServeEntity) -> ServerResponse: + """Convert the entity to a response + + Args: + entity (T): The entity + + Returns: + RES: The response + """ + # TODO implement your own logic here, transfer the entity to a response + return ServerResponse() diff --git a/kamidso/engine/serve/model/models/__init__.py b/kamidso/engine/serve/model/models/__init__.py new file mode 100644 index 0000000..d675f90 --- /dev/null +++ b/kamidso/engine/serve/model/models/__init__.py @@ -0,0 +1,2 @@ +# This is an auto-generated __init__.py file +# generated by `kamidso.engine.core new serve model` diff --git a/kamidso/engine/serve/model/models/model_adapter.py b/kamidso/engine/serve/model/models/model_adapter.py new file mode 100644 index 0000000..870bf91 --- /dev/null +++ b/kamidso/engine/serve/model/models/model_adapter.py @@ -0,0 +1,73 @@ +import json +from typing import Type + +from sqlalchemy.orm import Session + +from kamidso.engine.core.core.interface.storage import StorageItemAdapter +from kamidso.engine.core.model.cluster.storage import ModelStorageIdentifier, ModelStorageItem + +from .models import ServeEntity + + +class ModelStorageAdapter(StorageItemAdapter[ModelStorageItem, ServeEntity]): + """File metadata adapter. + + Convert between storage format and database model. + """ + + def to_storage_format(self, item: ModelStorageItem) -> ServeEntity: + """Convert to storage format.""" + params = json.dumps(item.params, ensure_ascii=False) + enabled = 1 if item.enabled else 0 + return ServeEntity( + host=item.host, + port=item.port, + model=item.model, + provider=item.provider, + worker_type=item.worker_type, + enabled=enabled, + worker_name=item.worker_name, + params=params, + description=item.description, + user_name=item.user_name, + sys_code=item.sys_code, + ) + + def from_storage_format(self, model: ServeEntity) -> ModelStorageItem: + """Convert from storage format.""" + params = json.loads(model.params) + enabled = True if model.enabled else False + return ModelStorageItem( + host=model.host, + port=model.port, + model=model.model, + provider=model.provider, + worker_type=model.worker_type, + enabled=enabled, + worker_name=model.worker_name, + params=params, + description=model.description, + user_name=model.user_name, + sys_code=model.sys_code, + ) + + def get_query_for_identifier( + self, + storage_format: Type[ServeEntity], + resource_id: ModelStorageIdentifier, + **kwargs, + ): + """Get query for identifier.""" + session: Session = kwargs.get("session") + if session is None: + raise Exception("session is None") + query = ( + session.query(storage_format) + .filter(storage_format.model == resource_id.model) + .filter(storage_format.worker_type == resource_id.worker_type) + ) + if resource_id.user_name: + query = query.filter(storage_format.user_name == resource_id.user_name) + if resource_id.sys_code: + query = query.filter(storage_format.sys_code == resource_id.sys_code) + return query diff --git a/kamidso/engine/serve/model/models/models.py b/kamidso/engine/serve/model/models/models.py new file mode 100644 index 0000000..76e492e --- /dev/null +++ b/kamidso/engine/serve/model/models/models.py @@ -0,0 +1,98 @@ +"""This is an auto-generated model file +You can define your own models and DAOs here +""" + +from datetime import datetime +from typing import Any, Dict, Union + +from sqlalchemy import Column, DateTime, Integer, String, Text, UniqueConstraint + +from kamidso.engine.core.storage.metadata import BaseDao, Model + +from ..api.schemas import ServeRequest, ServerResponse +from ..config import SERVER_APP_TABLE_NAME, ServeConfig + + +class ServeEntity(Model): + __tablename__ = SERVER_APP_TABLE_NAME + __table_args__ = ( + UniqueConstraint( + "model", "provider", "worker_type", name="uk_model_provider_type" + ), + ) + id = Column(Integer, primary_key=True, comment="Auto increment id") + + host = Column(String(255), nullable=False, comment="The model worker host") + port = Column(Integer, nullable=False, comment="The model worker port") + model = Column(String(255), nullable=False, comment="The model name") + provider = Column(String(255), nullable=False, comment="The model provider") + worker_type = Column(String(255), nullable=False, comment="The worker type") + params = Column(Text, nullable=False, comment="The model parameters, JSON format") + enabled = Column( + Integer, + default=1, + nullable=True, + comment="Whether the model is enabled, if it is enabled, it will be started " + "when the system starts, 1 is enabled, 0 is disabled", + ) + worker_name = Column(String(255), nullable=True, comment="The worker name") + description = Column(Text, nullable=True, comment="The model description") + user_name = Column(String(128), index=True, nullable=True, comment="User name") + sys_code = Column(String(128), index=True, nullable=True, comment="System code") + + gmt_created = Column(DateTime, default=datetime.now, comment="Record creation time") + gmt_modified = Column(DateTime, default=datetime.now, comment="Record update time") + + def __repr__(self): + return ( + f"ServeEntity(id={self.id}, gmt_created='{self.gmt_created}', " + f"gmt_modified='{self.gmt_modified}')" + ) + + +class ServeDao(BaseDao[ServeEntity, ServeRequest, ServerResponse]): + """The DAO class for Model""" + + def __init__(self, serve_config: ServeConfig): + super().__init__() + self._serve_config = serve_config + + def from_request(self, request: Union[ServeRequest, Dict[str, Any]]) -> ServeEntity: + """Convert the request to an entity + + Args: + request (Union[ServeRequest, Dict[str, Any]]): The request + + Returns: + T: The entity + """ + request_dict = ( + request.to_dict() if isinstance(request, ServeRequest) else request + ) + entity = ServeEntity(**request_dict) + # TODO implement your own logic here, transfer the request_dict to an entity + return entity + + def to_request(self, entity: ServeEntity) -> ServeRequest: + """Convert the entity to a request + + Args: + entity (T): The entity + + Returns: + REQ: The request + """ + # TODO implement your own logic here, transfer the entity to a request + return ServeRequest() + + def to_response(self, entity: ServeEntity) -> ServerResponse: + """Convert the entity to a response + + Args: + entity (T): The entity + + Returns: + RES: The response + """ + # TODO implement your own logic here, transfer the entity to a response + return ServerResponse() diff --git a/kamidso/engine/serve/prompt/models/__init__.py b/kamidso/engine/serve/prompt/models/__init__.py new file mode 100644 index 0000000..1d296f8 --- /dev/null +++ b/kamidso/engine/serve/prompt/models/__init__.py @@ -0,0 +1,2 @@ +# This is an auto-generated __init__.py file +# generated by `kamidso.engine.core new serve prompt` diff --git a/kamidso/engine/serve/prompt/models/models.py b/kamidso/engine/serve/prompt/models/models.py new file mode 100644 index 0000000..7b54993 --- /dev/null +++ b/kamidso/engine/serve/prompt/models/models.py @@ -0,0 +1,151 @@ +"""This is an auto-generated model file +You can define your own models and DAOs here +""" + +import uuid +from datetime import datetime +from typing import Any, Dict, Union + +from sqlalchemy import Column, DateTime, Integer, String, Text, UniqueConstraint + +from kamidso.engine.core.storage.metadata import BaseDao, Model + +from ..api.schemas import ServeRequest, ServerResponse +from ..config import ServeConfig + + +class ServeEntity(Model): + __tablename__ = "prompt_manage" + __table_args__ = ( + UniqueConstraint( + "prompt_name", + "sys_code", + "prompt_language", + "model", + name="uk_prompt_name_sys_code", + ), + ) + id = Column(Integer, primary_key=True, comment="Auto increment id") + + chat_scene = Column(String(100), comment="Chat scene") + sub_chat_scene = Column(String(100), comment="Sub chat scene") + prompt_code = Column(String(256), comment="Prompt Code") + prompt_type = Column(String(100), comment="Prompt type(eg: common, private)") + prompt_name = Column(String(256), comment="Prompt name") + content = Column(Text, comment="Prompt content") + input_variables = Column( + String(1024), nullable=True, comment="Prompt input variables(split by comma))" + ) + response_schema = Column(Text, comment="Prompt response schema") + model = Column( + String(128), + nullable=True, + comment="Prompt model name(we can use different models for different prompt", + ) + prompt_language = Column( + String(32), index=True, nullable=True, comment="Prompt language(eg:en, zh-cn)" + ) + prompt_format = Column( + String(32), + index=True, + nullable=True, + default="f-string", + comment="Prompt format(eg: f-string, jinja2)", + ) + prompt_desc = Column(String(512), nullable=True, comment="Prompt description") + user_code = Column(String(128), index=True, nullable=True, comment="User code") + user_name = Column(String(128), index=True, nullable=True, comment="User name") + sys_code = Column(String(128), index=True, nullable=True, comment="System code") + gmt_created = Column(DateTime, default=datetime.now, comment="Record creation time") + gmt_modified = Column(DateTime, default=datetime.now, comment="Record update time") + + def __repr__(self): + return ( + f"ServeEntity(id={self.id}, chat_scene='{self.chat_scene}', " + f"sub_chat_scene='{self.sub_chat_scene}', " + f"prompt_type='{self.prompt_type}', prompt_name='{self.prompt_name}', " + f"content='{self.content}'," + f"user_code='{self.user_code}', user_name='{self.user_name}', " + f"gmt_created='{self.gmt_created}', gmt_modified='{self.gmt_modified}')" + ) + + +class ServeDao(BaseDao[ServeEntity, ServeRequest, ServerResponse]): + """The DAO class for Prompt""" + + def __init__(self, serve_config: ServeConfig): + super().__init__() + self._serve_config = serve_config + + def from_request(self, request: Union[ServeRequest, Dict[str, Any]]) -> ServeEntity: + """Convert the request to an entity + + Args: + request (Union[ServeRequest, Dict[str, Any]]): The request + + Returns: + T: The entity + """ + request_dict = request.dict() if isinstance(request, ServeRequest) else request + entity = ServeEntity(**request_dict) + if not entity.prompt_code: + entity.prompt_code = uuid.uuid4().hex + return entity + + def to_request(self, entity: ServeEntity) -> ServeRequest: + """Convert the entity to a request + + Args: + entity (T): The entity + + Returns: + REQ: The request + """ + return ServeRequest( + chat_scene=entity.chat_scene, + sub_chat_scene=entity.sub_chat_scene, + prompt_type=entity.prompt_type, + prompt_name=entity.prompt_name, + prompt_code=entity.prompt_code, + content=entity.content, + response_schema=entity.response_schema, + input_variables=entity.input_variables, + model=entity.model, + prompt_language=entity.prompt_language, + prompt_desc=entity.prompt_desc, + user_code=entity.user_code, + user_name=entity.user_name, + sys_code=entity.sys_code, + ) + + def to_response(self, entity: ServeEntity) -> ServerResponse: + """Convert the entity to a response + + Args: + entity (T): The entity + + Returns: + RES: The response + """ + # TODO implement your own logic here, transfer the entity to a response + gmt_created_str = entity.gmt_created.strftime("%Y-%m-%d %H:%M:%S") + gmt_modified_str = entity.gmt_modified.strftime("%Y-%m-%d %H:%M:%S") + return ServerResponse( + id=entity.id, + chat_scene=entity.chat_scene, + sub_chat_scene=entity.sub_chat_scene, + prompt_code=entity.prompt_code, + prompt_type=entity.prompt_type, + prompt_name=entity.prompt_name, + content=entity.content, + prompt_desc=entity.prompt_desc, + user_name=entity.user_name, + user_code=entity.user_code, + model=entity.model, + input_variables=entity.input_variables, + response_schema=entity.response_schema, + prompt_language=entity.prompt_language, + sys_code=entity.sys_code, + gmt_created=gmt_created_str, + gmt_modified=gmt_modified_str, + ) diff --git a/kamidso/engine/serve/prompt/models/prompt_template_adapter.py b/kamidso/engine/serve/prompt/models/prompt_template_adapter.py new file mode 100644 index 0000000..04a60de --- /dev/null +++ b/kamidso/engine/serve/prompt/models/prompt_template_adapter.py @@ -0,0 +1,56 @@ +from typing import Type + +from sqlalchemy.orm import Session + +from kamidso.engine.core.core.interface.prompt import PromptTemplateIdentifier, StoragePromptTemplate +from kamidso.engine.core.core.interface.storage import StorageItemAdapter + +from .models import ServeEntity + + +class PromptTemplateAdapter(StorageItemAdapter[StoragePromptTemplate, ServeEntity]): + def to_storage_format(self, item: StoragePromptTemplate) -> ServeEntity: + return ServeEntity( + chat_scene=item.chat_scene, + sub_chat_scene=item.sub_chat_scene, + prompt_type=item.prompt_type, + prompt_name=item.prompt_name, + content=item.content, + input_variables=item.input_variables, + model=item.model, + prompt_language=item.prompt_language, + prompt_format=item.prompt_format, + user_name=item.user_name, + sys_code=item.sys_code, + ) + + def from_storage_format(self, model: ServeEntity) -> StoragePromptTemplate: + return StoragePromptTemplate( + chat_scene=model.chat_scene, + sub_chat_scene=model.sub_chat_scene, + prompt_type=model.prompt_type, + prompt_name=model.prompt_name, + content=model.content, + input_variables=model.input_variables, + model=model.model, + prompt_language=model.prompt_language, + prompt_format=model.prompt_format, + user_name=model.user_name, + sys_code=model.sys_code, + ) + + def get_query_for_identifier( + self, + storage_format: Type[ServeEntity], + resource_id: PromptTemplateIdentifier, + **kwargs, + ): + session: Session = kwargs.get("session") + if session is None: + raise Exception("session is None") + query_obj = session.query(ServeEntity) + for key, value in resource_id.to_dict().items(): + if value is None: + continue + query_obj = query_obj.filter(getattr(ServeEntity, key) == value) + return query_obj diff --git a/kamidso/engine/serve/rag/models/__init__.py b/kamidso/engine/serve/rag/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/kamidso/engine/serve/rag/models/chunk_db.py b/kamidso/engine/serve/rag/models/chunk_db.py new file mode 100644 index 0000000..bfea6a8 --- /dev/null +++ b/kamidso/engine/serve/rag/models/chunk_db.py @@ -0,0 +1,278 @@ +from datetime import datetime +from typing import Any, Dict, List, Union + +from sqlalchemy import Column, DateTime, Integer, String, Text, func, not_ + +from kamidso.engine.core._private.pydantic import model_to_dict +from kamidso.engine.core.storage.metadata import BaseDao, Model +from kamidso.engine.core.storage.metadata._base_dao import QUERY_SPEC, REQ, RES +from kamidso.engine.serve.rag.api.schemas import ChunkServeRequest, ChunkServeResponse + + +class DocumentChunkEntity(Model): + __tablename__ = "document_chunk" + id = Column(Integer, primary_key=True) + document_id = Column(Integer) + doc_name = Column(String(100)) + doc_type = Column(String(100)) + content = Column(Text) + questions = Column(Text) + meta_info = Column(String(500)) + gmt_created = Column(DateTime) + gmt_modified = Column(DateTime) + + def __repr__(self): + return ( + f"DocumentChunkEntity(id={self.id}, doc_name='{self.doc_name}', " + f"doc_type='{self.doc_type}', " + f"document_id='{self.document_id}', content='{self.content}', " + f"questions='{self.questions}', meta_info='{self.meta_info}', " + f"gmt_created='{self.gmt_created}', gmt_modified='{self.gmt_modified}')" + ) + + def to_dict(self): + return { + "id": self.id, + "document_id": self.document_id, + "doc_name": self.doc_name, + "doc_type": self.doc_type, + "content": self.content, + "questions": self.questions, + "meta_info": self.meta_info, + "gmt_created": self.gmt_created, + "gmt_modified": self.gmt_modified, + } + + +class DocumentChunkDao(BaseDao): + def create_documents_chunks(self, documents: List): + session = self.get_raw_session() + docs = [ + DocumentChunkEntity( + doc_name=document.doc_name, + doc_type=document.doc_type, + document_id=document.document_id, + content=document.content or "", + meta_info=document.meta_info or "", + gmt_created=datetime.now(), + gmt_modified=datetime.now(), + ) + for document in documents + ] + session.add_all(docs) + session.commit() + session.close() + + def get_document_chunks( + self, query: DocumentChunkEntity, page=1, page_size=20, document_ids=None + ): + session = self.get_raw_session() + document_chunks = session.query(DocumentChunkEntity) + if query.id is not None: + document_chunks = document_chunks.filter(DocumentChunkEntity.id == query.id) + if query.document_id is not None: + document_chunks = document_chunks.filter( + DocumentChunkEntity.document_id == query.document_id + ) + if query.doc_type is not None: + document_chunks = document_chunks.filter( + DocumentChunkEntity.doc_type == query.doc_type + ) + if query.content is not None: + document_chunks = document_chunks.filter( + DocumentChunkEntity.content.like(f"%{query.content}%") + ) + if query.doc_name is not None: + document_chunks = document_chunks.filter( + DocumentChunkEntity.doc_name == query.doc_name + ) + if query.meta_info is not None: + document_chunks = document_chunks.filter( + DocumentChunkEntity.meta_info == query.meta_info + ) + if document_ids is not None: + document_chunks = document_chunks.filter( + DocumentChunkEntity.document_id.in_(document_ids) + ) + + document_chunks = document_chunks.order_by(DocumentChunkEntity.id.asc()) + document_chunks = document_chunks.offset((page - 1) * page_size).limit( + page_size + ) + result = document_chunks.all() + session.close() + return result + + def get_chunks_with_questions(self, query: DocumentChunkEntity, document_ids=None): + session = self.get_raw_session() + document_chunks = session.query(DocumentChunkEntity) + if query.doc_name is not None: + document_chunks = document_chunks.filter( + DocumentChunkEntity.doc_name == query.doc_name + ) + if query.meta_info is not None: + document_chunks = document_chunks.filter( + DocumentChunkEntity.meta_info == query.meta_info + ) + document_chunks = document_chunks.filter( + not_(DocumentChunkEntity.questions is None) + ) + if document_ids is not None: + document_chunks = document_chunks.filter( + DocumentChunkEntity.document_id.in_(document_ids) + ) + + document_chunks = document_chunks.order_by(DocumentChunkEntity.id.asc()) + result = document_chunks.all() + session.close() + return result + + def update(self, query_request: QUERY_SPEC, update_request: REQ) -> RES: + """Update an entity object. + + Args: + query_request (REQ): The request schema object or dict for query. + update_request (REQ): The request schema object for update. + Returns: + RES: The response schema object. + """ + with self.session() as session: + query = self._create_query_object(session, query_request) + entry = query.first() + if entry is None: + raise Exception("Invalid request") + for key, value in model_to_dict(update_request).items(): # type: ignore + if value is not None: + if key in ["gmt_created", "gmt_modified"]: + # Assuming the datetime format is 'YYYY-MM-DD HH:MM:SS' + value = datetime.strptime(value, "%Y-%m-%d %H:%M:%S") + setattr(entry, key, value) + session.merge(entry) + return self.to_response(entry) + + def update_chunk(self, chunk: DocumentChunkEntity): + """Update a chunk""" + try: + session = self.get_raw_session() + updated = session.merge(chunk) + session.commit() + return updated.id + finally: + session.close() + + def get_document_chunks_count(self, query: DocumentChunkEntity): + session = self.get_raw_session() + document_chunks = session.query(func.count(DocumentChunkEntity.id)) + if query.id is not None: + document_chunks = document_chunks.filter(DocumentChunkEntity.id == query.id) + if query.document_id is not None: + document_chunks = document_chunks.filter( + DocumentChunkEntity.document_id == query.document_id + ) + if query.doc_type is not None: + document_chunks = document_chunks.filter( + DocumentChunkEntity.doc_type == query.doc_type + ) + if query.doc_name is not None: + document_chunks = document_chunks.filter( + DocumentChunkEntity.doc_name == query.doc_name + ) + if query.meta_info is not None: + document_chunks = document_chunks.filter( + DocumentChunkEntity.meta_info == query.meta_info + ) + count = document_chunks.scalar() + session.close() + return count + + def raw_delete(self, document_id: int): + session = self.get_raw_session() + if document_id is None: + raise Exception("document_id is None") + query = DocumentChunkEntity(document_id=document_id) + knowledge_documents = session.query(DocumentChunkEntity) + if query.document_id is not None: + chunks = knowledge_documents.filter( + DocumentChunkEntity.document_id == query.document_id + ) + chunks.delete() + session.commit() + session.close() + + def from_request( + self, request: Union[ChunkServeRequest, Dict[str, Any]] + ) -> DocumentChunkEntity: + """Convert the request to an entity + + Args: + request (Union[ServeRequest, Dict[str, Any]]): The request + + Returns: + T: The entity + """ + request_dict = ( + request.dict() if isinstance(request, ChunkServeRequest) else request + ) + entity = DocumentChunkEntity(**request_dict) + return entity + + def to_request(self, entity: DocumentChunkEntity) -> ChunkServeRequest: + """Convert the entity to a request + + Args: + entity (T): The entity + + Returns: + REQ: The request + """ + return ChunkServeRequest( + id=entity.id, + doc_name=entity.doc_name, + doc_type=entity.doc_type, + document_id=entity.document_id, + content=entity.content, + questions=entity.questions, + meta_info=entity.meta_info, + gmt_created=entity.gmt_created, + gmt_modified=entity.gmt_modified, + ) + + def to_response(self, entity: DocumentChunkEntity) -> ChunkServeResponse: + """Convert the entity to a response + + Args: + entity (T): The entity + + Returns: + REQ: The request + """ + gmt_created_str = entity.gmt_created.strftime("%Y-%m-%d %H:%M:%S") + gmt_modified_str = entity.gmt_modified.strftime("%Y-%m-%d %H:%M:%S") + return ChunkServeResponse( + id=entity.id, + doc_name=entity.doc_name, + doc_type=entity.doc_type, + document_id=entity.document_id, + content=entity.content, + questions=entity.questions, + meta_info=entity.meta_info, + gmt_created=gmt_created_str, + gmt_modified=gmt_modified_str, + ) + + def from_response( + self, response: Union[ChunkServeResponse, Dict[str, Any]] + ) -> DocumentChunkEntity: + """Convert the request to an entity + + Args: + request (Union[ServeRequest, Dict[str, Any]]): The request + + Returns: + T: The entity + """ + response_dict = ( + response.dict() if isinstance(response, ChunkServeResponse) else response + ) + entity = DocumentChunkEntity(**response_dict) + return entity diff --git a/kamidso/engine/serve/rag/models/code_graph_db.py b/kamidso/engine/serve/rag/models/code_graph_db.py new file mode 100644 index 0000000..d72feec --- /dev/null +++ b/kamidso/engine/serve/rag/models/code_graph_db.py @@ -0,0 +1,468 @@ +"""Code graph persistence models. + +Three tables: + - code_graph_vertex: AST-extracted code nodes (classes, functions, modules, etc.) + - code_graph_edge: Structural relationships (contains, calls, imports, etc.) + - code_graph_meta: Per-knowledge-space graph metadata (counts, build info) +""" + +import logging +from datetime import datetime +from typing import List, Optional + +from sqlalchemy import Column, DateTime, Integer, String, Text, func + +from kamidso.engine.core.storage.metadata import BaseDao, Model + +logger = logging.getLogger(__name__) + + +# Props that are extracted to dedicated columns (not stored in JSON props) +_VERTEX_PROPS_COLUMNS = {"node_type", "source_file", "language", "community"} + + +# --------------------------------------------------------------------------- +# Vertex +# --------------------------------------------------------------------------- + + +class CodeGraphVertexEntity(Model): + """Code graph vertex (node) entity.""" + + __tablename__ = "code_graph_vertex" + + id = Column(Integer, primary_key=True, autoincrement=True) + knowledge_id = Column(String(100), nullable=False) + vid = Column(String(500), nullable=False) + name = Column(String(500), nullable=False) + node_type = Column(String(50), nullable=False, default="") + source_file = Column(String(500), default="") + language = Column(String(30), default="") + community = Column(String(50), default="") + props = Column(Text, default=None) + gmt_create = Column(DateTime, default=datetime.now) + gmt_modified = Column(DateTime, default=datetime.now, onupdate=datetime.now) + + def __repr__(self): + return ( + f"CodeGraphVertexEntity(id={self.id}, vid='{self.vid}', " + f"name='{self.name}', node_type='{self.node_type}')" + ) + + def to_dict(self): + return { + "id": self.id, + "knowledge_id": self.knowledge_id, + "vid": self.vid, + "name": self.name, + "node_type": self.node_type, + "source_file": self.source_file, + "language": self.language, + "community": self.community, + "props": self.props, + } + + +class CodeGraphVertexDao(BaseDao): + """DAO for code_graph_vertex table.""" + + def batch_insert(self, knowledge_id: str, vertices: list) -> int: + """Bulk insert vertices for a knowledge space.""" + if not vertices: + return 0 + session = self.get_raw_session() + try: + current_time = datetime.now() + batch_size = 200 + for i in range(0, len(vertices), batch_size): + batch = vertices[i : i + batch_size] + mappings = [ + { + "knowledge_id": knowledge_id, + "vid": v["vid"], + "name": v["name"], + "node_type": v.get("node_type", ""), + "source_file": v.get("source_file", ""), + "language": v.get("language", ""), + "community": v.get("community", ""), + "props": v.get("props"), + "gmt_create": current_time, + "gmt_modified": current_time, + } + for v in batch + ] + session.bulk_insert_mappings(CodeGraphVertexEntity, mappings) + session.commit() + return len(vertices) + except Exception: + session.rollback() + raise + finally: + session.close() + + def get_by_knowledge_id(self, knowledge_id: str) -> List[CodeGraphVertexEntity]: + """Get all vertices for a knowledge space.""" + session = self.get_raw_session() + try: + return ( + session.query(CodeGraphVertexEntity) + .filter(CodeGraphVertexEntity.knowledge_id == knowledge_id) + .all() + ) + finally: + session.close() + + def search_by_name( + self, + knowledge_id: str, + name_prefix: str, + node_type: Optional[str] = None, + limit: int = 20, + ) -> List[CodeGraphVertexEntity]: + """Search vertices by name prefix (uses index).""" + session = self.get_raw_session() + try: + query = session.query(CodeGraphVertexEntity).filter( + CodeGraphVertexEntity.knowledge_id == knowledge_id, + CodeGraphVertexEntity.name.like(f"{name_prefix}%"), + ) + if node_type: + query = query.filter(CodeGraphVertexEntity.node_type == node_type) + return query.limit(limit).all() + finally: + session.close() + + def get_by_source_file( + self, knowledge_id: str, source_file: str + ) -> List[CodeGraphVertexEntity]: + """Get all vertices in a specific source file.""" + session = self.get_raw_session() + try: + return ( + session.query(CodeGraphVertexEntity) + .filter( + CodeGraphVertexEntity.knowledge_id == knowledge_id, + CodeGraphVertexEntity.source_file == source_file, + ) + .all() + ) + finally: + session.close() + + def get_by_vid( + self, knowledge_id: str, vid: str + ) -> Optional[CodeGraphVertexEntity]: + """Get a single vertex by its vid.""" + session = self.get_raw_session() + try: + return ( + session.query(CodeGraphVertexEntity) + .filter( + CodeGraphVertexEntity.knowledge_id == knowledge_id, + CodeGraphVertexEntity.vid == vid, + ) + .first() + ) + finally: + session.close() + + def delete_by_knowledge_id(self, knowledge_id: str) -> int: + """Delete all vertices for a knowledge space. Returns count deleted.""" + session = self.get_raw_session() + try: + count = ( + session.query(CodeGraphVertexEntity) + .filter(CodeGraphVertexEntity.knowledge_id == knowledge_id) + .delete() + ) + session.commit() + return count + except Exception: + session.rollback() + raise + finally: + session.close() + + def count_by_knowledge_id(self, knowledge_id: str) -> int: + """Count vertices for a knowledge space.""" + session = self.get_raw_session() + try: + result = ( + session.query(func.count(CodeGraphVertexEntity.id)) + .filter(CodeGraphVertexEntity.knowledge_id == knowledge_id) + .scalar() + ) + return result or 0 + finally: + session.close() + + +# --------------------------------------------------------------------------- +# Edge +# --------------------------------------------------------------------------- + + +class CodeGraphEdgeEntity(Model): + """Code graph edge (relationship) entity.""" + + __tablename__ = "code_graph_edge" + + id = Column(Integer, primary_key=True, autoincrement=True) + knowledge_id = Column(String(100), nullable=False) + sid = Column(String(500), nullable=False) + tid = Column(String(500), nullable=False) + edge_type = Column(String(50), nullable=False, default="references") + confidence = Column(String(20), default="EXTRACTED") + source_file = Column(String(500), default="") + source_location = Column(String(30), default="") + props = Column(Text, default=None) + gmt_create = Column(DateTime, default=datetime.now) + gmt_modified = Column(DateTime, default=datetime.now, onupdate=datetime.now) + + def __repr__(self): + return ( + f"CodeGraphEdgeEntity(id={self.id}, sid='{self.sid}', " + f"tid='{self.tid}', edge_type='{self.edge_type}')" + ) + + def to_dict(self): + return { + "id": self.id, + "knowledge_id": self.knowledge_id, + "sid": self.sid, + "tid": self.tid, + "edge_type": self.edge_type, + "confidence": self.confidence, + "source_file": self.source_file, + "source_location": self.source_location, + "props": self.props, + } + + +class CodeGraphEdgeDao(BaseDao): + """DAO for code_graph_edge table.""" + + def batch_insert(self, knowledge_id: str, edges: list) -> int: + """Bulk insert edges for a knowledge space.""" + if not edges: + return 0 + session = self.get_raw_session() + try: + current_time = datetime.now() + batch_size = 200 + for i in range(0, len(edges), batch_size): + batch = edges[i : i + batch_size] + mappings = [ + { + "knowledge_id": knowledge_id, + "sid": e["sid"], + "tid": e["tid"], + "edge_type": e["edge_type"], + "confidence": e.get("confidence", "EXTRACTED"), + "source_file": e.get("source_file", ""), + "source_location": e.get("source_location", ""), + "props": e.get("props"), + "gmt_create": current_time, + "gmt_modified": current_time, + } + for e in batch + ] + session.bulk_insert_mappings(CodeGraphEdgeEntity, mappings) + session.commit() + return len(edges) + except Exception: + session.rollback() + raise + finally: + session.close() + + def get_by_knowledge_id(self, knowledge_id: str) -> List[CodeGraphEdgeEntity]: + """Get all edges for a knowledge space.""" + session = self.get_raw_session() + try: + return ( + session.query(CodeGraphEdgeEntity) + .filter(CodeGraphEdgeEntity.knowledge_id == knowledge_id) + .all() + ) + finally: + session.close() + + def get_out_edges( + self, knowledge_id: str, sid: str, edge_type: Optional[str] = None + ) -> List[CodeGraphEdgeEntity]: + """Get outgoing edges from a vertex (sid matches).""" + session = self.get_raw_session() + try: + query = session.query(CodeGraphEdgeEntity).filter( + CodeGraphEdgeEntity.knowledge_id == knowledge_id, + CodeGraphEdgeEntity.sid == sid, + ) + if edge_type: + query = query.filter(CodeGraphEdgeEntity.edge_type == edge_type) + return query.all() + finally: + session.close() + + def get_in_edges( + self, knowledge_id: str, tid: str, edge_type: Optional[str] = None + ) -> List[CodeGraphEdgeEntity]: + """Get incoming edges to a vertex (tid matches).""" + session = self.get_raw_session() + try: + query = session.query(CodeGraphEdgeEntity).filter( + CodeGraphEdgeEntity.knowledge_id == knowledge_id, + CodeGraphEdgeEntity.tid == tid, + ) + if edge_type: + query = query.filter(CodeGraphEdgeEntity.edge_type == edge_type) + return query.all() + finally: + session.close() + + def delete_by_knowledge_id(self, knowledge_id: str) -> int: + """Delete all edges for a knowledge space. Returns count deleted.""" + session = self.get_raw_session() + try: + count = ( + session.query(CodeGraphEdgeEntity) + .filter(CodeGraphEdgeEntity.knowledge_id == knowledge_id) + .delete() + ) + session.commit() + return count + except Exception: + session.rollback() + raise + finally: + session.close() + + def count_by_knowledge_id(self, knowledge_id: str) -> int: + """Count edges for a knowledge space.""" + session = self.get_raw_session() + try: + result = ( + session.query(func.count(CodeGraphEdgeEntity.id)) + .filter(CodeGraphEdgeEntity.knowledge_id == knowledge_id) + .scalar() + ) + return result or 0 + finally: + session.close() + + +# --------------------------------------------------------------------------- +# Meta +# --------------------------------------------------------------------------- + + +class CodeGraphMetaEntity(Model): + """Code graph metadata entity — one row per knowledge space.""" + + __tablename__ = "code_graph_meta" + + id = Column(Integer, primary_key=True, autoincrement=True) + knowledge_id = Column(String(100), nullable=False, unique=True) + vertex_count = Column(Integer, default=0) + edge_count = Column(Integer, default=0) + community_count = Column(Integer, default=0) + build_source = Column(String(20), default="") + repo_url = Column(String(500), default="") + branch = Column(String(100), default="") + build_status = Column(String(20), default="completed") + graph_version = Column(Integer, default=1) + gmt_create = Column(DateTime, default=datetime.now) + gmt_modified = Column(DateTime, default=datetime.now, onupdate=datetime.now) + + def __repr__(self): + return ( + f"CodeGraphMetaEntity(id={self.id}, knowledge_id='{self.knowledge_id}', " + f"vertices={self.vertex_count}, edges={self.edge_count})" + ) + + def to_dict(self): + return { + "id": self.id, + "knowledge_id": self.knowledge_id, + "vertex_count": self.vertex_count, + "edge_count": self.edge_count, + "community_count": self.community_count, + "build_source": self.build_source, + "repo_url": self.repo_url, + "branch": self.branch, + "build_status": self.build_status, + "graph_version": self.graph_version, + "gmt_create": self.gmt_create.isoformat() if self.gmt_create else None, + "gmt_modified": ( + self.gmt_modified.isoformat() if self.gmt_modified else None + ), + } + + +class CodeGraphMetaDao(BaseDao): + """DAO for code_graph_meta table.""" + + def get_by_knowledge_id(self, knowledge_id: str) -> Optional[CodeGraphMetaEntity]: + """Get meta for a knowledge space.""" + session = self.get_raw_session() + try: + return ( + session.query(CodeGraphMetaEntity) + .filter(CodeGraphMetaEntity.knowledge_id == knowledge_id) + .first() + ) + finally: + session.close() + + def upsert(self, meta: CodeGraphMetaEntity) -> int: + """Insert or update meta for a knowledge space. + + Returns the row ID. + """ + session = self.get_raw_session() + try: + existing = ( + session.query(CodeGraphMetaEntity) + .filter(CodeGraphMetaEntity.knowledge_id == meta.knowledge_id) + .first() + ) + if existing: + existing.vertex_count = meta.vertex_count + existing.edge_count = meta.edge_count + existing.community_count = meta.community_count + existing.build_source = meta.build_source + existing.repo_url = meta.repo_url + existing.branch = meta.branch + existing.build_status = meta.build_status + existing.graph_version = meta.graph_version + existing.gmt_modified = datetime.now() + session.commit() + return existing.id + else: + meta.gmt_create = datetime.now() + meta.gmt_modified = datetime.now() + session.add(meta) + session.commit() + return meta.id + except Exception: + session.rollback() + raise + finally: + session.close() + + def delete_by_knowledge_id(self, knowledge_id: str) -> bool: + """Delete meta for a knowledge space.""" + session = self.get_raw_session() + try: + count = ( + session.query(CodeGraphMetaEntity) + .filter(CodeGraphMetaEntity.knowledge_id == knowledge_id) + .delete() + ) + session.commit() + return count > 0 + except Exception: + session.rollback() + raise + finally: + session.close() diff --git a/kamidso/engine/serve/rag/models/document_db.py b/kamidso/engine/serve/rag/models/document_db.py new file mode 100644 index 0000000..d8ba937 --- /dev/null +++ b/kamidso/engine/serve/rag/models/document_db.py @@ -0,0 +1,479 @@ +from datetime import datetime +from typing import Any, Dict, List, Union + +from sqlalchemy import Column, DateTime, Integer, String, Text, func + +from kamidso.engine.core._private.config import Config +from kamidso.engine.core._private.pydantic import model_to_dict +from kamidso.engine.core.storage.metadata import BaseDao, Model +from kamidso.engine.core.storage.metadata._base_dao import QUERY_SPEC, REQ, RES +from kamidso.engine.core.util import PaginationResult +from kamidso.engine.serve.conversation.api.schemas import ServeRequest +from kamidso.engine.serve.rag.api.schemas import ( + DocumentServeRequest, + DocumentServeResponse, +) + +CFG = Config() + + +class KnowledgeDocumentEntity(Model): + __tablename__ = "knowledge_document" + id = Column(Integer, primary_key=True) + doc_name = Column(String(100)) + doc_type = Column(String(100)) + doc_token = Column(String(100)) + space = Column(String(100)) + chunk_size = Column(Integer) + status = Column(String(100)) + last_sync = Column(DateTime) + content = Column(Text) + result = Column(Text) + vector_ids = Column(Text) + summary = Column(Text) + gmt_created = Column(DateTime) + gmt_modified = Column(DateTime) + questions = Column(Text) + + def __repr__(self): + return ( + f"KnowledgeDocumentEntity(id={self.id}, doc_name='{self.doc_name}', " + f"doc_type='{self.doc_type}', chunk_size='{self.chunk_size}', " + f"status='{self.status}', last_sync='{self.last_sync}', " + f"content='{self.content}', " + f"result='{self.result}', summary='{self.summary}', " + f"gmt_created='{self.gmt_created}', gmt_modified='{self.gmt_modified}', " + f"questions='{self.questions}')" + ) + + def to_dict(self): + return { + "__tablename__": self.__tablename__, + "id": self.id, + "doc_name": self.doc_name, + "doc_type": self.doc_type, + "doc_token": self.doc_token, + "space": self.space, + "chunk_size": self.chunk_size, + "status": self.status, + "last_sync": self.last_sync, + "content": self.content, + "result": self.result, + "vector_ids": self.vector_ids, + "summary": self.summary, + "gmt_create": self.gmt_created, + "gmt_modified": self.gmt_modified, + "questions": self.questions, + } + + +class KnowledgeDocumentDao(BaseDao): + def create_knowledge_document(self, document: KnowledgeDocumentEntity): + session = self.get_raw_session() + knowledge_document = KnowledgeDocumentEntity( + doc_name=document.doc_name, + doc_type=document.doc_type, + doc_token=document.doc_token, + space=document.space, + chunk_size=0.0, + status=document.status, + last_sync=document.last_sync, + content=document.content or "", + result=document.result or "", + summary=document.summary or "", + vector_ids=document.vector_ids, + gmt_created=datetime.now(), + gmt_modified=datetime.now(), + questions=document.questions, + ) + session.add(knowledge_document) + session.commit() + doc_id = knowledge_document.id + session.close() + return doc_id + + def get_knowledge_documents(self, query, page=1, page_size=20): + """Get a list of documents that match the given query. + Args: + query: A KnowledgeDocumentEntity object containing the query parameters. + page: The page number to return. + page_size: The number of documents to return per page. + """ + session = self.get_raw_session() + print(f"current session:{session}") + knowledge_documents = session.query(KnowledgeDocumentEntity) + if query.id is not None: + knowledge_documents = knowledge_documents.filter( + KnowledgeDocumentEntity.id == query.id + ) + if query.doc_name is not None: + knowledge_documents = knowledge_documents.filter( + KnowledgeDocumentEntity.doc_name == query.doc_name + ) + if query.doc_type is not None: + knowledge_documents = knowledge_documents.filter( + KnowledgeDocumentEntity.doc_type == query.doc_type + ) + if query.space is not None: + knowledge_documents = knowledge_documents.filter( + KnowledgeDocumentEntity.space == query.space + ) + if query.status is not None: + knowledge_documents = knowledge_documents.filter( + KnowledgeDocumentEntity.status == query.status + ) + + knowledge_documents = knowledge_documents.order_by( + KnowledgeDocumentEntity.id.desc() + ) + knowledge_documents = knowledge_documents.offset((page - 1) * page_size).limit( + page_size + ) + result = knowledge_documents.all() + session.close() + return result + + def document_by_id(self, document_id) -> KnowledgeDocumentEntity: + session = self.get_raw_session() + query = session.query(KnowledgeDocumentEntity).filter( + KnowledgeDocumentEntity.id == document_id + ) + + result = query.first() + session.close() + return result + + def documents_by_ids(self, ids) -> List[KnowledgeDocumentEntity]: + """Get a list of documents by their IDs. + Args: + ids: A list of document IDs. + Returns: + A list of KnowledgeDocumentEntity objects. + """ + session = self.get_raw_session() + print(f"current session:{session}") + knowledge_documents = session.query(KnowledgeDocumentEntity) + knowledge_documents = knowledge_documents.filter( + KnowledgeDocumentEntity.id.in_(ids) + ) + result = knowledge_documents.all() + session.close() + return result + + def get_documents(self, query): + session = self.get_raw_session() + print(f"current session:{session}") + knowledge_documents = session.query(KnowledgeDocumentEntity) + if query.id is not None: + knowledge_documents = knowledge_documents.filter( + KnowledgeDocumentEntity.id == query.id + ) + if query.doc_name is not None: + knowledge_documents = knowledge_documents.filter( + KnowledgeDocumentEntity.doc_name == query.doc_name + ) + if query.doc_type is not None: + knowledge_documents = knowledge_documents.filter( + KnowledgeDocumentEntity.doc_type == query.doc_type + ) + if query.space is not None: + knowledge_documents = knowledge_documents.filter( + KnowledgeDocumentEntity.space == query.space + ) + if query.status is not None: + knowledge_documents = knowledge_documents.filter( + KnowledgeDocumentEntity.status == query.status + ) + + knowledge_documents = knowledge_documents.order_by( + KnowledgeDocumentEntity.id.desc() + ) + result = knowledge_documents.all() + session.close() + return result + + def get_knowledge_documents_count_bulk(self, space_names): + session = self.get_raw_session() + """ + Perform a batch query to count the number of documents for each knowledge space. + + Args: + space_names: A list of knowledge space names to query for document counts. + session: A SQLAlchemy session object. + + Returns: + A dictionary mapping each space name to its document count. + """ + counts_query = ( + session.query( + KnowledgeDocumentEntity.space, + func.count(KnowledgeDocumentEntity.id).label("document_count"), + ) + .filter(KnowledgeDocumentEntity.space.in_(space_names)) + .group_by(KnowledgeDocumentEntity.space) + ) + + results = counts_query.all() + docs_count = {result.space: result.document_count for result in results} + session.close() + return docs_count + + def get_knowledge_documents_count_bulk_by_ids(self, spaces): + session = self.get_raw_session() + """ + Perform a batch query to count the number of documents for each knowledge space. + + Args: + spaces: A list of knowledge space names to query for document counts. + session: A SQLAlchemy session object. + + Returns: + A dictionary mapping each space name to its document count. + """ + # build the group by query + counts_query = ( + session.query( + KnowledgeDocumentEntity.space, + func.count(KnowledgeDocumentEntity.id).label("document_count"), + ) + .filter(KnowledgeDocumentEntity.space.in_(spaces)) + .group_by(KnowledgeDocumentEntity.space) + ) + + results = counts_query.all() + docs_count = {result.space: result.document_count for result in results} + session.close() + return docs_count + + def get_knowledge_documents_count(self, query): + session = self.get_raw_session() + knowledge_documents = session.query(func.count(KnowledgeDocumentEntity.id)) + if query.id is not None: + knowledge_documents = knowledge_documents.filter( + KnowledgeDocumentEntity.id == query.id + ) + if query.doc_name is not None: + knowledge_documents = knowledge_documents.filter( + KnowledgeDocumentEntity.doc_name == query.doc_name + ) + if query.doc_type is not None: + knowledge_documents = knowledge_documents.filter( + KnowledgeDocumentEntity.doc_type == query.doc_type + ) + if query.space is not None: + knowledge_documents = knowledge_documents.filter( + KnowledgeDocumentEntity.space == query.space + ) + if query.status is not None: + knowledge_documents = knowledge_documents.filter( + KnowledgeDocumentEntity.status == query.status + ) + count = knowledge_documents.scalar() + session.close() + return count + + def update(self, query_request: QUERY_SPEC, update_request: REQ) -> RES: + """Update an entity object. + + Args: + query_request (REQ): The request schema object or dict for query. + update_request (REQ): The request schema object for update. + Returns: + RES: The response schema object. + """ + with self.session() as session: + query = self._create_query_object(session, query_request) + entry = query.first() + if entry is None: + raise Exception("Invalid request") + for key, value in model_to_dict(update_request).items(): # type: ignore + if value is not None: + if key in ["last_sync", "gmt_created", "gmt_modified"]: + continue + setattr(entry, key, value) + session.merge(entry) + return self.to_response(entry) + + def update_set_space_id(self, space, space_id): + session = self.get_raw_session() + knowledge_documents = session.query(KnowledgeDocumentEntity) + if space is not None: + knowledge_documents.filter(KnowledgeDocumentEntity.space == space).filter( + KnowledgeDocumentEntity.id is None + ).update({KnowledgeDocumentEntity.id: space_id}, synchronize_session=False) + session.commit() + session.close() + + def update_knowledge_document(self, document: KnowledgeDocumentEntity): + try: + session = self.get_raw_session() + updated_document = session.merge(document) + session.commit() + return updated_document.id + finally: + session.close() + + def raw_delete(self, query: KnowledgeDocumentEntity): + session = self.get_raw_session() + knowledge_documents = session.query(KnowledgeDocumentEntity) + if query.id is not None: + knowledge_documents = knowledge_documents.filter( + KnowledgeDocumentEntity.id == query.id + ) + if query.doc_name is not None: + knowledge_documents = knowledge_documents.filter( + KnowledgeDocumentEntity.doc_name == query.doc_name + ) + if query.space is not None: + knowledge_documents = knowledge_documents.filter( + KnowledgeDocumentEntity.space == query.space + ) + knowledge_documents.delete() + session.commit() + session.close() + + def get_list_page( + self, query_request: QUERY_SPEC, page: int, page_size: int + ) -> PaginationResult[RES]: + """Get a page of entity objects. + + Args: + query_request (REQ): The request schema object or dict for query. + page (int): The page number. + page_size (int): The page size. + + Returns: + PaginationResult: The pagination result. + """ + with self.session() as session: + query = self._create_query_object(session, query_request) + total_count = query.count() + items = ( + query.order_by(KnowledgeDocumentEntity.id.desc()) + .offset((page - 1) * page_size) + .limit(page_size) + ) + items = [self.to_response(item) for item in items] + total_pages = (total_count + page_size - 1) // page_size + + return PaginationResult( + items=items, + total_count=total_count, + total_pages=total_pages, + page=page, + page_size=page_size, + ) + + def from_request( + self, request: Union[ServeRequest, Dict[str, Any]] + ) -> KnowledgeDocumentEntity: + """Convert the request to an entity + + Args: + request (Union[ServeRequest, Dict[str, Any]]): The request + + Returns: + T: The entity + """ + request_dict = ( + request.dict() if isinstance(request, DocumentServeRequest) else request + ) + entity = KnowledgeDocumentEntity( + id=request_dict.get("id"), + doc_name=request_dict.get("doc_name"), + doc_type=request_dict.get("doc_type"), + doc_token=request_dict.get("doc_token"), + space=request_dict.get("space"), + chunk_size=request_dict.get("chunk_size"), + status=request_dict.get("status"), + last_sync=request_dict.get("last_sync"), + content=request_dict.get("content"), + result=request_dict.get("result"), + vector_ids=request_dict.get("vector_ids"), + summary=request_dict.get("summary"), + questions=request_dict.get("questions"), + ) + return entity + + def to_request(self, entity: KnowledgeDocumentEntity) -> DocumentServeResponse: + """Convert the entity to a request + + Args: + entity (T): The entity + + Returns: + REQ: The request + """ + return DocumentServeResponse( + id=entity.id, + doc_name=entity.doc_name, + doc_type=entity.doc_type, + space=entity.space, + chunk_size=entity.chunk_size, + status=entity.status, + last_sync=entity.last_sync, + content=entity.content, + result=entity.result, + vector_ids=entity.vector_ids, + summary=entity.summary, + questions=entity.questions, + gmt_created=entity.gmt_created, + gmt_modified=entity.gmt_modified, + ) + + def to_response(self, entity: KnowledgeDocumentEntity) -> DocumentServeResponse: + """Convert the entity to a response + + Args: + entity (T): The entity + + Returns: + REQ: The request + """ + return DocumentServeResponse( + id=entity.id, + doc_name=entity.doc_name, + doc_type=entity.doc_type, + space=entity.space, + chunk_size=entity.chunk_size, + status=entity.status, + last_sync=str(entity.last_sync), + content=entity.content, + result=entity.result, + vector_ids=entity.vector_ids, + summary=entity.summary, + questions=entity.questions, + gmt_created=str(entity.gmt_created), + gmt_modified=str(entity.gmt_modified), + ) + + def from_response( + self, response: Union[DocumentServeResponse, Dict[str, Any]] + ) -> KnowledgeDocumentEntity: + """Convert the request to an entity + + Args: + request (Union[ServeRequest, Dict[str, Any]]): The request + + Returns: + T: The entity + """ + response_dict = ( + response.dict() if isinstance(response, DocumentServeResponse) else response + ) + entity = KnowledgeDocumentEntity( + id=response_dict.get("id"), + doc_name=response_dict.get("doc_name"), + doc_type=response_dict.get("doc_type"), + doc_token=response_dict.get("doc_token"), + space=response_dict.get("space"), + chunk_size=response_dict.get("chunk_size"), + status=response_dict.get("status"), + last_sync=response_dict.get("last_sync"), + content=response_dict.get("content"), + result=response_dict.get("result"), + vector_ids=response_dict.get("vector_ids"), + summary=response_dict.get("summary"), + questions=response_dict.get("questions"), + ) + return entity diff --git a/kamidso/engine/serve/rag/models/models.py b/kamidso/engine/serve/rag/models/models.py new file mode 100644 index 0000000..fceb668 --- /dev/null +++ b/kamidso/engine/serve/rag/models/models.py @@ -0,0 +1,197 @@ +from datetime import datetime +from typing import Any, Dict, Union + +from sqlalchemy import Column, DateTime, Integer, String, Text + +from kamidso.engine.core._private.pydantic import model_to_dict +from kamidso.engine.core.storage.metadata import BaseDao, Model +from kamidso.engine.app.knowledge.request.request import KnowledgeSpaceRequest +from kamidso.engine.serve.rag.api.schemas import SpaceServeRequest, SpaceServeResponse + + +class KnowledgeSpaceEntity(Model): + __tablename__ = "knowledge_space" + id = Column(Integer, primary_key=True) + name = Column(String(100)) + vector_type = Column(String(100)) + index_methods = Column(String(500)) # JSON string of index methods + domain_type = Column(String(100)) + desc = Column(String(100)) + owner = Column(String(100)) + context = Column(Text) + gmt_created = Column(DateTime) + gmt_modified = Column(DateTime) + + def __repr__(self): + return ( + f"KnowledgeSpaceEntity(id={self.id}, name='{self.name}', " + f"vector_type='{self.vector_type}', desc='{self.desc}', " + f"owner='{self.owner}' context='{self.context}', " + f"gmt_created='{self.gmt_created}', gmt_modified='{self.gmt_modified}')" + ) + + +class KnowledgeSpaceDao(BaseDao): + def create_knowledge_space(self, space: KnowledgeSpaceRequest): + """Create knowledge space""" + import json + + session = self.get_raw_session() + # Convert index_methods list to JSON string + index_methods_json = None + if space.index_methods: + index_methods_json = json.dumps(space.index_methods) + knowledge_space = KnowledgeSpaceEntity( + name=space.name, + vector_type=space.vector_type, + index_methods=index_methods_json, + domain_type=space.domain_type, + desc=space.desc, + owner=space.owner, + gmt_created=datetime.now(), + gmt_modified=datetime.now(), + ) + session.add(knowledge_space) + session.commit() + session.close() + ks = self.get_knowledge_space(KnowledgeSpaceEntity(name=space.name)) + if ks is not None and len(ks) == 1: + return ks[0].id + raise Exception("create space error, find more than 1 or 0 space.") + + def get_knowledge_space_by_ids(self, ids): + session = self.get_raw_session() + if ids: + knowledge_spaces = session.query(KnowledgeSpaceEntity).filter( + KnowledgeSpaceEntity.id.in_(ids) + ) + else: + return [] + knowledge_spaces_list = knowledge_spaces.all() + session.close() + return knowledge_spaces_list + + def get_knowledge_space(self, query: KnowledgeSpaceEntity): + """Get knowledge space by query""" + session = self.get_raw_session() + knowledge_spaces = session.query(KnowledgeSpaceEntity) + if query.id is not None: + knowledge_spaces = knowledge_spaces.filter( + KnowledgeSpaceEntity.id == query.id + ) + if query.name is not None: + knowledge_spaces = knowledge_spaces.filter( + KnowledgeSpaceEntity.name == query.name + ) + if query.vector_type is not None: + knowledge_spaces = knowledge_spaces.filter( + KnowledgeSpaceEntity.vector_type == query.vector_type + ) + if query.domain_type is not None: + knowledge_spaces = knowledge_spaces.filter( + KnowledgeSpaceEntity.domain_type == query.domain_type + ) + if query.desc is not None: + knowledge_spaces = knowledge_spaces.filter( + KnowledgeSpaceEntity.desc == query.desc + ) + if query.owner is not None: + knowledge_spaces = knowledge_spaces.filter( + KnowledgeSpaceEntity.owner == query.owner + ) + if query.gmt_created is not None: + knowledge_spaces = knowledge_spaces.filter( + KnowledgeSpaceEntity.gmt_created == query.gmt_created + ) + if query.gmt_modified is not None: + knowledge_spaces = knowledge_spaces.filter( + KnowledgeSpaceEntity.gmt_modified == query.gmt_modified + ) + knowledge_spaces = knowledge_spaces.order_by( + KnowledgeSpaceEntity.gmt_created.desc() + ) + result = knowledge_spaces.all() + session.close() + return result + + def update_knowledge_space(self, space: KnowledgeSpaceEntity): + """Update knowledge space""" + + session = self.get_raw_session() + request = SpaceServeRequest(id=space.id) + update_request = self.to_request(space) + query = self._create_query_object(session, request) + entry = query.first() + if entry is None: + raise Exception("Invalid request") + for key, value in model_to_dict(update_request).items(): # type: ignore + if value is not None: + setattr(entry, key, value) + session.merge(entry) + session.commit() + session.close() + return self.to_response(space) + + def delete_knowledge_space(self, space: KnowledgeSpaceEntity): + """Delete knowledge space""" + session = self.get_raw_session() + if space: + session.delete(space) + session.commit() + session.close() + + def from_request( + self, request: Union[SpaceServeRequest, Dict[str, Any]] + ) -> KnowledgeSpaceEntity: + """Convert the request to an entity + + Args: + request (Union[ServeRequest, Dict[str, Any]]): The request + + Returns: + T: The entity + """ + request_dict = ( + model_to_dict(request) + if isinstance(request, SpaceServeRequest) + else request + ) + entity = KnowledgeSpaceEntity(**request_dict) + return entity + + def to_request(self, entity: KnowledgeSpaceEntity) -> SpaceServeRequest: + """Convert the entity to a request + + Args: + entity (T): The entity + + Returns: + REQ: The request + """ + return SpaceServeRequest( + id=entity.id, + name=entity.name, + vector_type=entity.vector_type, + desc=entity.desc, + owner=entity.owner, + context=entity.context, + ) + + def to_response(self, entity: KnowledgeSpaceEntity) -> SpaceServeResponse: + """Convert the entity to a response + + Args: + entity (T): The entity + + Returns: + REQ: The request + """ + return SpaceServeResponse( + id=entity.id, + name=entity.name, + vector_type=entity.vector_type, + desc=entity.desc, + owner=entity.owner, + context=entity.context, + domain_type=entity.domain_type, + ) diff --git a/kamidso/engine/serve/session_file/models/__init__.py b/kamidso/engine/serve/session_file/models/__init__.py new file mode 100644 index 0000000..ac49df7 --- /dev/null +++ b/kamidso/engine/serve/session_file/models/__init__.py @@ -0,0 +1,6 @@ +"""Persistence models for session-scoped files.""" + +from .dao import SessionFileDao +from .models import SessionFileEntity + +__all__ = ["SessionFileDao", "SessionFileEntity"] diff --git a/kamidso/engine/serve/session_file/models/dao.py b/kamidso/engine/serve/session_file/models/dao.py new file mode 100644 index 0000000..a4240db --- /dev/null +++ b/kamidso/engine/serve/session_file/models/dao.py @@ -0,0 +1,253 @@ +"""Repository operations for session-scoped files.""" + +import re +from typing import Any, Dict, List, Optional, Union + +from sqlalchemy import func + +from kamidso.engine.core.storage.metadata import BaseDao + +from ..domain import ( + FileScope, + SessionFileManifest, + SessionFilePrivateRecord, + SessionFileSnapshot, + SessionFileStatus, +) +from .models import SessionFileEntity + + +class SessionFileDao(BaseDao[SessionFileEntity, Dict[str, Any], SessionFileManifest]): + """Store private file locations while exposing only public metadata.""" + + def from_request(self, request: Union[Dict[str, Any], Any]) -> SessionFileEntity: + if isinstance(request, dict): + request_dict = request + elif hasattr(request, "model_dump"): + request_dict = request.model_dump() + else: + request_dict = request.dict() + return SessionFileEntity(**request_dict) + + def create(self, request: Union[Dict[str, Any], Any]) -> SessionFileManifest: + """Validate scope and insert without using unrestricted BaseDao lookup.""" + entity = self.from_request(request) + self._validate_entity(entity) + with self.session() as session: + session.add(entity) + session.flush() + return self.to_manifest(entity) + + def to_response(self, entity: SessionFileEntity) -> SessionFileManifest: + """Map fields explicitly so private storage never crosses this boundary.""" + return self.to_manifest(entity) + + @staticmethod + def to_manifest(entity: SessionFileEntity) -> SessionFileManifest: + """Build agent-safe file metadata from a persistence record.""" + return SessionFileManifest( + file_id=entity.file_id, + name=entity.display_name, + size=entity.size_bytes, + media_type=entity.media_type, + kind=entity.file_kind, + status=SessionFileStatus(entity.status), + ordinal=entity.ordinal, + ) + + @staticmethod + def to_snapshot(entity: SessionFileEntity) -> SessionFileSnapshot: + """Build history-safe file metadata from a persistence record.""" + return SessionFileSnapshot( + file_id=entity.file_id, + name=entity.display_name, + size=entity.size_bytes, + media_type=entity.media_type, + kind=entity.file_kind, + status=SessionFileStatus(entity.status), + ordinal=entity.ordinal, + ) + + @staticmethod + def to_private_record(entity: SessionFileEntity) -> SessionFilePrivateRecord: + """Detach complete private metadata into an immutable domain value.""" + return SessionFilePrivateRecord( + file_id=entity.file_id, + owner_id=entity.owner_id, + session_id=entity.session_id, + task_id=entity.task_id, + display_name=entity.display_name, + storage_uri=entity.storage_uri, + media_type=entity.media_type, + file_kind=entity.file_kind, + size_bytes=entity.size_bytes, + sha256=entity.sha256, + ordinal=entity.ordinal, + status=SessionFileStatus(entity.status), + inspection_json=entity.inspection_json, + error_code=entity.error_code, + error_message=entity.error_message, + source_file_id=entity.source_file_id, + created_at=entity.created_at, + updated_at=entity.updated_at, + ) + + def get_by_file_id( + self, file_id: str, scope: FileScope + ) -> Optional[SessionFileManifest]: + """Return a file only when owner and scope match, otherwise ``None``.""" + with self.session(commit=False) as session: + entity = ( + self._scope_query(session, scope) + .filter(SessionFileEntity.file_id == file_id) + .first() + ) + return self.to_manifest(entity) if entity is not None else None + + def get_private_file_by_id( + self, file_id: str, scope: FileScope + ) -> Optional[SessionFilePrivateRecord]: + """Return immutable private metadata for an exact owner-bound scope.""" + with self.session(commit=False) as session: + entity = ( + self._scope_query(session, scope) + .filter(SessionFileEntity.file_id == file_id) + .first() + ) + return self.to_private_record(entity) if entity is not None else None + + def list_by_scope(self, scope: FileScope) -> List[SessionFileManifest]: + """List an owner's scoped files in stable upload order.""" + with self.session(commit=False) as session: + entities = ( + self._scope_query(session, scope) + .order_by(SessionFileEntity.ordinal) + .all() + ) + return [self.to_manifest(entity) for entity in entities] + + def list_by_source_file_id( + self, source_file_id: str, scope: FileScope + ) -> List[SessionFileManifest]: + """List lineage records within one exact owner-bound scope.""" + with self.session(commit=False) as session: + entities = ( + self._scope_query(session, scope) + .filter(SessionFileEntity.source_file_id == source_file_id) + .order_by(SessionFileEntity.ordinal) + .all() + ) + return [self.to_manifest(entity) for entity in entities] + + def total_size_bytes(self, scope: FileScope) -> int: + """Return total persisted bytes for one owner-bound scope.""" + with self.session(commit=False) as session: + total = ( + self._scope_query(session, scope) + .with_entities(func.coalesce(func.sum(SessionFileEntity.size_bytes), 0)) + .scalar() + ) + return int(total) + + def total_owner_size_bytes(self, owner_id: str) -> int: + """Return total persisted bytes across all scopes for one owner.""" + if not isinstance(owner_id, str) or not owner_id.strip(): + raise ValueError("owner_id must not be blank") + with self.session(commit=False) as session: + total = ( + session.query(func.coalesce(func.sum(SessionFileEntity.size_bytes), 0)) + .filter(SessionFileEntity.owner_id == owner_id) + .scalar() + ) + return int(total) + + def update_status( + self, + file_id: str, + scope: FileScope, + status: Union[SessionFileStatus, str], + inspection_json: Optional[str] = None, + error_code: Optional[str] = None, + error_message: Optional[str] = None, + ) -> Optional[SessionFileManifest]: + """Update processing state only for the exact owner-bound scope.""" + status_value = self._validate_status(status) + with self.session() as session: + entity = ( + self._scope_query(session, scope) + .filter(SessionFileEntity.file_id == file_id) + .first() + ) + if entity is None: + return None + entity.status = status_value + entity.inspection_json = inspection_json + entity.error_code = error_code + entity.error_message = error_message + session.flush() + return self.to_manifest(entity) + + def delete_by_file_id(self, file_id: str, scope: FileScope) -> bool: + """Delete a file only for the exact owner-bound scope.""" + with self.session() as session: + deleted = ( + self._scope_query(session, scope) + .filter(SessionFileEntity.file_id == file_id) + .delete(synchronize_session=False) + ) + return deleted == 1 + + def get_one(self, *args, **kwargs): + raise NotImplementedError("Use a scope-constrained session file method") + + def get_list(self, *args, **kwargs): + raise NotImplementedError("Use a scope-constrained session file method") + + def get_list_page(self, *args, **kwargs): + raise NotImplementedError("Use a scope-constrained session file method") + + def update(self, *args, **kwargs): + raise NotImplementedError("Use a scope-constrained session file method") + + def delete(self, *args, **kwargs): + raise NotImplementedError("Use a scope-constrained session file method") + + @staticmethod + def _validate_entity(entity: SessionFileEntity) -> None: + FileScope( + owner_id=entity.owner_id, + session_id=entity.session_id, + task_id=entity.task_id, + ) + SessionFileDao._validate_status(entity.status) + if entity.size_bytes < 0: + raise ValueError("size_bytes must be non-negative") + if entity.ordinal < 0: + raise ValueError("ordinal must be non-negative") + if ( + not isinstance(entity.sha256, str) + or re.fullmatch(r"[0-9a-fA-F]{64}", entity.sha256) is None + ): + raise ValueError("sha256 must contain exactly 64 hexadecimal characters") + + @staticmethod + def _validate_status(status: Union[SessionFileStatus, str]) -> str: + try: + return SessionFileStatus(status).value + except (TypeError, ValueError) as error: + raise ValueError("invalid session file status") from error + + @staticmethod + def _scope_query(session, scope: FileScope): + query = session.query(SessionFileEntity).filter( + SessionFileEntity.owner_id == scope.owner_id + ) + if scope.session_id is not None: + return query.filter( + SessionFileEntity.session_id == scope.session_id, + SessionFileEntity.task_id.is_(None), + ) + return query.filter( + SessionFileEntity.task_id == scope.task_id, + SessionFileEntity.session_id.is_(None), + ) diff --git a/kamidso/engine/serve/session_file/models/models.py b/kamidso/engine/serve/session_file/models/models.py new file mode 100644 index 0000000..4d7303f --- /dev/null +++ b/kamidso/engine/serve/session_file/models/models.py @@ -0,0 +1,81 @@ +"""SQLAlchemy model for session-scoped files.""" + +from datetime import datetime + +from sqlalchemy import ( + BigInteger, + CheckConstraint, + Column, + DateTime, + Index, + Integer, + String, + Text, + func, +) +from sqlalchemy.dialects import mysql + +from kamidso.engine.core.storage.metadata import Model + + +class SessionFileEntity(Model): + """Private persistence record for an owner-bound session or task file.""" + + __tablename__ = "dbgpt_session_file" + __table_args__ = ( + CheckConstraint( + "(session_id IS NULL) <> (task_id IS NULL)", + name="ck_session_file_scope", + ), + Index("uk_session_file_file_id", "file_id", unique=True), + Index( + "idx_session_file_owner_session", + "owner_id", + "session_id", + "ordinal", + ), + Index( + "idx_session_file_owner_task", + "owner_id", + "task_id", + "ordinal", + ), + Index("idx_session_file_sha256", "owner_id", "sha256"), + ) + + id = Column( + BigInteger().with_variant(Integer, "sqlite"), + primary_key=True, + autoincrement=True, + ) + file_id = Column(String(64), nullable=False) + owner_id = Column(String(255), nullable=False) + session_id = Column(String(255), nullable=True) + task_id = Column(String(64), nullable=True) + display_name = Column(String(256), nullable=False) + storage_uri = Column(String(512), nullable=False) + media_type = Column(String(255), nullable=False) + file_kind = Column(String(32), nullable=False) + size_bytes = Column(BigInteger, nullable=False) + sha256 = Column(String(64), nullable=False) + ordinal = Column(Integer, nullable=False) + status = Column(String(32), nullable=False) + inspection_json = Column( + Text().with_variant(mysql.LONGTEXT(), "mysql"), nullable=True + ) + error_code = Column(String(64), nullable=True) + error_message = Column(Text, nullable=True) + source_file_id = Column(String(64), nullable=True) + created_at = Column( + DateTime, + nullable=False, + default=datetime.now, + server_default=func.now(), + ) + updated_at = Column( + DateTime, + nullable=False, + default=datetime.now, + onupdate=datetime.now, + server_default=func.now(), + ) diff --git a/kamidso/engine/serve/utils/_template_files/default_serve_template/models/__init__.py b/kamidso/engine/serve/utils/_template_files/default_serve_template/models/__init__.py new file mode 100644 index 0000000..7a1b7a9 --- /dev/null +++ b/kamidso/engine/serve/utils/_template_files/default_serve_template/models/__init__.py @@ -0,0 +1,2 @@ +# This is an auto-generated __init__.py file +# generated by `kamidso.engine.core new serve {__template_app_name__}` diff --git a/kamidso/engine/serve/utils/_template_files/default_serve_template/models/models.py b/kamidso/engine/serve/utils/_template_files/default_serve_template/models/models.py new file mode 100644 index 0000000..13241b8 --- /dev/null +++ b/kamidso/engine/serve/utils/_template_files/default_serve_template/models/models.py @@ -0,0 +1,77 @@ +"""This is an auto-generated model file +You can define your own models and DAOs here +""" + +from datetime import datetime +from typing import Any, Dict, Union + +from sqlalchemy import Column, DateTime, Integer + +from kamidso.engine.core.storage.metadata import BaseDao, Model + +from ..api.schemas import ServeRequest, ServerResponse +from ..config import SERVER_APP_TABLE_NAME, ServeConfig + + +class ServeEntity(Model): + __tablename__ = SERVER_APP_TABLE_NAME + id = Column(Integer, primary_key=True, comment="Auto increment id") + + # TODO: define your own fields here + + gmt_created = Column(DateTime, default=datetime.now, comment="Record creation time") + gmt_modified = Column(DateTime, default=datetime.now, comment="Record update time") + + def __repr__(self): + return ( + f"ServeEntity(id={self.id}, gmt_created='{self.gmt_created}', " + f"gmt_modified='{self.gmt_modified}')" + ) + + +class ServeDao(BaseDao[ServeEntity, ServeRequest, ServerResponse]): + """The DAO class for {__template_app_name__hump__}""" + + def __init__(self, serve_config: ServeConfig): + super().__init__() + self._serve_config = serve_config + + def from_request(self, request: Union[ServeRequest, Dict[str, Any]]) -> ServeEntity: + """Convert the request to an entity + + Args: + request (Union[ServeRequest, Dict[str, Any]]): The request + + Returns: + T: The entity + """ + request_dict = ( + request.to_dict() if isinstance(request, ServeRequest) else request + ) + entity = ServeEntity(**request_dict) + # TODO implement your own logic here, transfer the request_dict to an entity + return entity + + def to_request(self, entity: ServeEntity) -> ServeRequest: + """Convert the entity to a request + + Args: + entity (T): The entity + + Returns: + REQ: The request + """ + # TODO implement your own logic here, transfer the entity to a request + return ServeRequest() + + def to_response(self, entity: ServeEntity) -> ServerResponse: + """Convert the entity to a response + + Args: + entity (T): The entity + + Returns: + RES: The response + """ + # TODO implement your own logic here, transfer the entity to a response + return ServerResponse()