feat:custom dataset

Author: LuckyYC

backend/config/__init__.py

@@ -9,7 +9,7 @@
 from sqlalchemy.orm import declarative_base
 from sqlalchemy.pool import AsyncAdaptedQueuePool
 
-from config.db_config import get_settings
+from db.db_config import get_settings
 
 settings = get_settings()
 

backend/config/db_config.py backend/db/db_config.pybackend/config/db_config.py renamed to backend/db/db_config.py

@@ -160,15 +160,15 @@ class TaskCreateReq(BaseModel):
     system_prompt: Optional[str] = Field(
         default="", description="System prompt for the model"
     )
-    user_prompt: Optional[str] = Field(
-        default="", description="User prompt for the model"
-    )
     request_payload: Optional[str] = Field(
         default="", description="Custom request payload for non-chat APIs (JSON string)"
     )
     field_mapping: Optional[Dict[str, str]] = Field(
         default=None, description="Field mapping configuration for custom APIs"
     )
+    test_data: Optional[str] = Field(
+        default="", description="Custom test data in JSONL format or file path"
+    )
 
 
 class TaskResultItem(BaseModel):
@@ -315,7 +315,6 @@ class Task(Base):
     target_host = Column(String(255), nullable=False)
     model = Column(String(100), nullable=True)
     system_prompt = Column(Text, nullable=True)
-    user_prompt = Column(Text, nullable=True)
     stream_mode = Column(String(20), nullable=False)
     concurrent_users = Column(Integer, nullable=False)
     spawn_rate = Column(Integer, nullable=False)
@@ -331,6 +330,7 @@ class Task(Base):
     request_payload = Column(Text, nullable=True)
     field_mapping = Column(Text, nullable=True)
     error_message = Column(Text, nullable=True)
+    test_data = Column(Text, nullable=True)
     created_at = Column(DateTime, server_default=func.now())
     updated_at = Column(DateTime, server_default=func.now(), onupdate=func.now())
 

backend/db/mysql.py

@@ -32,12 +32,14 @@ class UploadFileRsp(BaseModel):
         task_id: The identifier for the associated task.
         files: A list of `UploadedFileInfo` objects for each uploaded file.
         cert_config: An optional dictionary for certificate configuration.
+        test_data: An optional string for test data file path.
     """
 
     message: str
     task_id: str
     files: List[UploadedFileInfo]
     cert_config: Optional[dict] = None
+    test_data: Optional[str] = None
 
 
 class UploadFileReq(BaseModel):

backend/model/task.py

@@ -32,24 +32,34 @@ known_third_party = ["fastapi", "pydantic", "sqlalchemy", "loguru", "uvicorn"]
 
 [tool.mypy]
 python_version = "3.9"
+# completely disable type checking
+ignore_errors = true
+ignore_missing_imports = true
+follow_imports = "skip"
 # base settings
 warn_return_any = false
-warn_unused_configs = true
-show_error_codes = true
+warn_unused_configs = false
+show_error_codes = false
 
 # relax type check requirements
 disallow_untyped_defs = false
 disallow_incomplete_defs = false
 check_untyped_defs = false
 disallow_untyped_decorators = false
 no_implicit_optional = false
+disallow_any_generics = false
+disallow_any_unimported = false
+disallow_subclassing_any = false
 
-# keep useful warnings
+# disable warnings
 warn_redundant_casts = false
 warn_unused_ignores = false
 warn_no_return = false
 warn_unreachable = false
 strict_equality = false
+strict_optional = false
+allow_redefinition = true
+implicit_reexport = true
 
 # ignore common third-party libraries
 [[tool.mypy.overrides]]
@@ -82,6 +92,14 @@ disallow_incomplete_defs = false
 check_untyped_defs = false
 warn_return_any = false
 ignore_errors = true
+follow_imports = "skip"
+follow_imports_for_stubs = false
+
+# SQLAlchemy specific overrides
+[[tool.mypy.overrides]]
+module = ["sqlalchemy.*"]
+ignore_missing_imports = true
+ignore_errors = true
 
 [tool.pytest.ini_options]
 testpaths = ["tests"]

backend/model/upload.py

@@ -7,9 +7,9 @@
 
 from starlette.responses import JSONResponse
 
-from config.config import LOG_DIR
 from model.log import LogContentResponse
-from utils.logger import be_logger as logger
+from utils.be_config import LOG_DIR
+from utils.logger import logger
 
 
 def get_last_n_lines(file_path: str, n: int = 100) -> str:
@@ -36,15 +36,15 @@ def get_last_n_lines(file_path: str, n: int = 100) -> str:
                 all_lines = f.readlines()
                 return "".join(all_lines[-n:]) if all_lines else ""
 
-            # For larger files, use a more reliable approach
-            # Start from end and read backwards in larger chunks
-            buffer_size = 8192
-            lines: list[str] = []
+            # For larger files, use a more efficient approach
+            # Read from end in chunks and collect lines
+            lines = list[str]()
             buffer = ""
             position = file_size
+            buffer_size = 8192
 
             while position > 0 and len(lines) < n:
-                # Calculate chunk size to read
+                # Calculate how much to read
                 chunk_size = min(buffer_size, position)
                 position -= chunk_size
 
@@ -56,26 +56,45 @@ def get_last_n_lines(file_path: str, n: int = 100) -> str:
                 buffer = chunk + buffer
 
                 # Split buffer into lines
-                lines_in_buffer = buffer.split("\n")
-
-                # Keep the first part (might be incomplete line) in buffer
-                buffer = lines_in_buffer[0]
-
-                # Add complete lines to our lines list (in reverse order since we're reading backwards)
-                for line in reversed(lines_in_buffer[1:]):
+                split_lines = buffer.split("\n")
+
+                # If we haven't reached the beginning, keep the first part as incomplete line
+                if position > 0:
+                    buffer = split_lines[0]
+                    # Process complete lines (skip the first incomplete one)
+                    complete_lines = split_lines[1:]
+                else:
+                    # At the beginning of file, all lines are complete
+                    buffer = ""
+                    complete_lines = split_lines
+
+                # Add complete lines to the front of our lines list
+                # (since we're reading backwards)
+                for line in reversed(complete_lines):
                     lines.insert(0, line)
                     if len(lines) >= n:
                         break
 
-                # If we've reached the beginning of file, add the remaining buffer as a line
-                if position == 0 and buffer:
-                    lines.insert(0, buffer)
+                # If we have enough lines, break
+                if len(lines) >= n:
+                    break
 
-            # Take last n lines and join them with newlines
+            # Take last n lines
             result_lines = lines[-n:] if len(lines) > n else lines
-            return "\n".join(result_lines) + (
-                "\n" if result_lines and not result_lines[-1].endswith("\n") else ""
-            )
+
+            # Join lines and ensure proper ending
+            if not result_lines:
+                return ""
+
+            result = "\n".join(result_lines)
+            # Add final newline if the original content had one and result doesn't end with one
+            if result and not result.endswith("\n"):
+                # Check if original file ends with newline
+                f.seek(max(0, file_size - 1))
+                if f.read(1) == "\n":
+                    result += "\n"
+
+            return result
 
     except Exception as e:
         logger.error(f"Failed to read log file: {str(e)}")

backend/pyproject.toml

@@ -12,7 +12,6 @@
 from sqlalchemy import func, or_, select, text
 from starlette.responses import JSONResponse
 
-from config.config import UPLOAD_FOLDER
 from model.task import (
     ComparisonMetrics,
     ComparisonRequest,
@@ -29,7 +28,8 @@
     TaskResultRsp,
     TaskStatusRsp,
 )
-from utils.logger import be_logger as logger
+from utils.be_config import UPLOAD_FOLDER
+from utils.logger import logger
 
 
 async def get_tasks_svc(
@@ -53,7 +53,7 @@ async def get_tasks_svc(
     Returns:
         A `TaskResponse` object containing the list of tasks and pagination details.
     """
-    task_list = []
+    task_list: List[Dict] = []
     pagination = Pagination()
     try:
         db = request.state.db
@@ -100,27 +100,67 @@ async def get_tasks_svc(
         )
 
         # Format the task data for the response.
-        task_list = [
-            {
+        task_list = []
+        for task in tasks:
+            # Convert headers from JSON string back to a list of objects for the frontend.
+            # headers_list = []
+            # if task.headers:
+            #     try:
+            #         headers_dict = json.loads(task.headers)
+            #         headers_list = [
+            #             {"key": k, "value": v} for k, v in headers_dict.items()
+            #         ]
+            #     except json.JSONDecodeError:
+            #         logger.warning(
+            #             f"Could not parse headers JSON for task {task.id}: {task.headers}"
+            #         )
+
+            # Convert cookies from JSON string back to a list of objects for the frontend.
+            # cookies_list = []
+            # if task.cookies:
+            #     try:
+            #         cookies_dict = json.loads(task.cookies)
+            #         cookies_list = [
+            #             {"key": k, "value": v} for k, v in cookies_dict.items()
+            #         ]
+            #     except json.JSONDecodeError:
+            #         logger.warning(
+            #             f"Could not parse cookies JSON for task {task.id}: {task.cookies}"
+            #         )
+
+            # Parse field_mapping from JSON string back to dictionary
+            field_mapping_dict = {}
+            if task.field_mapping:
+                try:
+                    field_mapping_dict = json.loads(task.field_mapping)
+                except json.JSONDecodeError:
+                    logger.warning(
+                        f"Could not parse field_mapping JSON for task {task.id}: {task.field_mapping}"
+                    )
+
+            task_data = {
                 "id": task.id,
                 "name": task.name,
                 "status": task.status,
                 "target_host": task.target_host,
                 "api_path": task.api_path,
                 "model": task.model,
                 "request_payload": task.request_payload,
-                "field_mapping": task.field_mapping,
+                "field_mapping": field_mapping_dict,
                 "concurrent_users": task.concurrent_users,
                 "duration": task.duration,
                 "spawn_rate": task.spawn_rate,
                 "chat_type": task.chat_type,
                 "stream_mode": str(task.stream_mode).lower() == "true",
-                "error_message": task.error_message,
+                "headers": "",
+                "cookies": "",
+                "cert_config": "",
+                "system_prompt": task.system_prompt or "",
+                "test_data": task.test_data or "",
                 "created_at": task.created_at.isoformat() if task.created_at else None,
                 "updated_at": task.updated_at.isoformat() if task.updated_at else None,
             }
-            for task in tasks
-        ]
+            task_list.append(task_data)
     except Exception as e:
         logger.error(f"Error getting tasks: {e}", exc_info=True)
         return TaskResponse(data=[], pagination=Pagination(), status="error")
@@ -227,17 +267,19 @@ async def create_task_svc(request: Request, body: TaskCreateReq):
     # Convert absolute paths to relative paths for cross-service compatibility
     if cert_file:
         # Convert backend upload path to relative path that st_engine can access
-        if cert_file.startswith("/app/upload_files/"):
-            cert_file = cert_file.replace("/app/upload_files/", "")
-        elif UPLOAD_FOLDER in cert_file:
+        if cert_file.startswith(UPLOAD_FOLDER + "/"):
             cert_file = cert_file.replace(UPLOAD_FOLDER + "/", "")
+        elif cert_file.startswith("/app/upload_files/"):
+            # For backward compatibility with existing Docker paths
+            cert_file = cert_file.replace("/app/upload_files/", "")
 
     if key_file:
         # Convert backend upload path to relative path that st_engine can access
-        if key_file.startswith("/app/upload_files/"):
-            key_file = key_file.replace("/app/upload_files/", "")
-        elif UPLOAD_FOLDER in key_file:
+        if key_file.startswith(UPLOAD_FOLDER + "/"):
             key_file = key_file.replace(UPLOAD_FOLDER + "/", "")
+        elif key_file.startswith("/app/upload_files/"):
+            # For backward compatibility with existing Docker paths
+            key_file = key_file.replace("/app/upload_files/", "")
 
     # Convert headers from a list of objects to a dictionary, then to a JSON string.
     headers = {
@@ -278,12 +320,12 @@ async def create_task_svc(request: Request, body: TaskCreateReq):
             status="created",
             error_message="",
             system_prompt=body.system_prompt,
-            user_prompt=body.user_prompt,
             cert_file=cert_file,
             key_file=key_file,
             api_path=body.api_path,
             request_payload=body.request_payload,
             field_mapping=field_mapping_json,
+            test_data=body.test_data,
         )
 
         db.add(new_task)
@@ -433,11 +475,11 @@ async def get_task_svc(request: Request, task_id: str):
             "headers": headers_list,
             "cookies": cookies_list,
             "cert_config": {"cert_file": task.cert_file, "key_file": task.key_file},
-            "system_prompt": task.system_prompt,
-            "user_prompt": task.user_prompt,
+            "system_prompt": task.system_prompt or "",
             "api_path": task.api_path,
             "request_payload": task.request_payload,
             "field_mapping": field_mapping_dict,
+            "test_data": task.test_data or "",
             "error_message": task.error_message,
             "created_at": task.created_at.isoformat() if task.created_at else None,
             "updated_at": task.updated_at.isoformat() if task.updated_at else None,
@@ -738,11 +780,12 @@ def _prepare_request_payload(body: TaskCreateReq) -> Dict:
     """Prepare request payload based on API path and configuration."""
     if body.api_path == "/v1/chat/completions":
         # Use the traditional chat completions format
-        messages = []
-        if body.system_prompt:
-            messages.append({"role": "system", "content": body.system_prompt})
-
-        messages.append({"role": "user", "content": body.user_prompt or "Hi"})
+        messages = [
+            {
+                "role": "user",
+                "content": "Hi",
+            }
+        ]
 
         return {
             "model": body.model,
@@ -931,7 +974,7 @@ async def _handle_streaming_response(response, full_url: str) -> Dict:
             }
 
         # For testing purposes, we limit the time and data we collect
-        max_chunks = 200  # max chunks to collect for testing
+        max_chunks = 300  # max chunks to collect for testing
         max_duration = 15  # max duration to wait for testing
 
         start_time = asyncio.get_event_loop().time()