Merge pull request #2 from MigoXLab/feat/custom_api

Feat/custom api

Author: LuckyYC

backend/api/api_task.py

@@ -27,6 +27,7 @@
     get_tasks_status_svc,
     get_tasks_svc,
     stop_task_svc,
+    test_api_endpoint_svc,
 )
 
 # Create an API router for task-related endpoints
@@ -175,3 +176,18 @@ async def compare_performance(request: Request, comparison_request: ComparisonRe
         ComparisonResponse: A response object containing comparison metrics.
     """
     return await compare_performance_svc(request, comparison_request)
+
+
+@router.post("/test", response_model=Dict[str, Any])
+async def test_api_endpoint(request: Request, task_create: TaskCreateReq):
+    """
+    Test the API endpoint with the provided configuration.
+
+    Args:
+        request (Request): The incoming request.
+        task_create (TaskCreateReq): The data for testing the API endpoint.
+
+    Returns:
+        Dict[str, Any]: A response containing the test result.
+    """
+    return await test_api_endpoint_svc(request, task_create)

backend/db/mysql.py

@@ -5,8 +5,8 @@
 
 from typing import AsyncGenerator
 
-from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
-from sqlalchemy.orm import declarative_base, sessionmaker
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
+from sqlalchemy.orm import declarative_base
 from sqlalchemy.pool import AsyncAdaptedQueuePool
 
 from config.db_config import get_settings
@@ -28,7 +28,7 @@
 )
 
 # Create a factory for asynchronous database sessions
-async_session_factory = sessionmaker(
+async_session_factory = async_sessionmaker(
     engine,
     class_=AsyncSession,
     expire_on_commit=False,

backend/model/task.py

@@ -88,6 +88,19 @@ class HeaderItem(BaseModel):
     fixed: bool = True
 
 
+class CookieItem(BaseModel):
+    """
+    Represents a single HTTP cookie item for a request.
+
+    Attributes:
+        key: The cookie name.
+        value: The cookie value.
+    """
+
+    key: str
+    value: str
+
+
 class CertConfig(BaseModel):
     """
     Configuration for SSL/TLS certificates.
@@ -123,19 +136,24 @@ class TaskCreateReq(BaseModel):
     api_path: str = Field(
         default="/v1/chat/completions", description="API path to test"
     )
-    model: str = Field(..., description="Name of the model to test")
+    model: Optional[str] = Field(default="", description="Name of the model to test")
     duration: int = Field(
         default=300, ge=1, description="Duration of the test in seconds"
     )
     concurrent_users: int = Field(..., ge=1, description="Number of concurrent users")
     spawn_rate: int = Field(ge=1, description="Number of users to spawn per second")
-    chat_type: int = Field(ge=0, description="Type of chat interaction")
+    chat_type: Optional[int] = Field(
+        default=0, ge=0, description="Type of chat interaction"
+    )
     stream_mode: bool = Field(
         default=True, description="Whether to use streaming response"
     )
     headers: List[HeaderItem] = Field(
         default_factory=list, description="List of request headers"
     )
+    cookies: List[CookieItem] = Field(
+        default_factory=list, description="List of request cookies"
+    )
     cert_config: Optional[CertConfig] = Field(
         default=None, description="Certificate configuration"
     )
@@ -145,6 +163,12 @@ class TaskCreateReq(BaseModel):
     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"
+    )
 
 
 class TaskResultItem(BaseModel):
@@ -289,20 +313,23 @@ class Task(Base):
     name = Column(String(255), nullable=False)
     status = Column(String(32), nullable=False)
     target_host = Column(String(255), nullable=False)
-    model = Column(String(100), 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)
     duration = Column(Integer, nullable=False)
-    chat_type = Column(Integer, nullable=False)
+    chat_type = Column(Integer, nullable=True)
     log_file = Column(Text, nullable=True)
     result_file = Column(Text, nullable=True)
     cert_file = Column(String(255), nullable=True)
     key_file = Column(String(255), nullable=True)
     headers = Column(Text, nullable=True)
-    # api_path = Column(String(255), nullable=True)
+    cookies = Column(Text, nullable=True)
+    api_path = Column(String(255), nullable=True)
+    request_payload = Column(Text, nullable=True)
+    field_mapping = Column(Text, nullable=True)
     error_message = Column(Text, nullable=True)
     created_at = Column(DateTime, server_default=func.now())
     updated_at = Column(DateTime, server_default=func.now(), onupdate=func.now())
@@ -326,10 +353,10 @@ class TaskResult(Base):
     p90_latency = Column(Float, nullable=False)
     rps = Column(Float, nullable=False)
     avg_content_length = Column(Float, nullable=False)
-    total_tps = Column(Float, nullable=False)
-    completion_tps = Column(Float, nullable=False)
-    avg_total_tokens_per_req = Column(Float, nullable=False)
-    avg_completion_tokens_per_req = Column(Float, nullable=False)
+    total_tps = Column(Float, nullable=True, default=0.0)
+    completion_tps = Column(Float, nullable=True, default=0.0)
+    avg_total_tokens_per_req = Column(Float, nullable=True, default=0.0)
+    avg_completion_tokens_per_req = Column(Float, nullable=True, default=0.0)
     created_at = Column(DateTime, server_default=func.now())
     updated_at = Column(DateTime, server_default=func.now(), onupdate=func.now())
 
@@ -349,8 +376,18 @@ def to_task_result_item(self) -> TaskResultItem:
             rps=self.rps,
             avg_content_length=self.avg_content_length,
             created_at=self.created_at.isoformat() if self.created_at else "",
-            total_tps=self.total_tps,
-            completion_tps=self.completion_tps,
-            avg_total_tokens_per_req=self.avg_total_tokens_per_req,
-            avg_completion_tokens_per_req=self.avg_completion_tokens_per_req,
+            total_tps=self.total_tps if self.total_tps is not None else 0.0,
+            completion_tps=(
+                self.completion_tps if self.completion_tps is not None else 0.0
+            ),
+            avg_total_tokens_per_req=(
+                self.avg_total_tokens_per_req
+                if self.avg_total_tokens_per_req is not None
+                else 0.0
+            ),
+            avg_completion_tokens_per_req=(
+                self.avg_completion_tokens_per_req
+                if self.avg_completion_tokens_per_req is not None
+                else 0.0
+            ),
         )

backend/mypy.ini

@@ -0,0 +1,32 @@
+[mypy]
+python_version = 3.11
+warn_return_any = False
+warn_unused_configs = True
+show_error_codes = True
+
+# ignore SQLAlchemy errors
+disallow_untyped_defs = False
+disallow_incomplete_defs = False
+check_untyped_defs = False
+
+# ignore SQLAlchemy module import errors
+[mypy-sqlalchemy.*]
+ignore_missing_imports = True
+ignore_errors = True
+
+# ignore specific SQLAlchemy type errors
+[mypy-model.*]
+ignore_errors = True
+
+# ignore other dependencies
+[mypy-fastapi.*]
+ignore_missing_imports = True
+
+[mypy-pydantic.*]
+ignore_missing_imports = True
+
+[mypy-uvicorn.*]
+ignore_missing_imports = True
+
+[mypy-starlette.*]
+ignore_missing_imports = True

backend/pyproject.toml

@@ -103,4 +103,4 @@ markers = [
 
 [tool.bandit]
 exclude_dirs = ["tests"]
-skips = ["B101", "B601"]
+skips = ["B101", "B601", "B501"]

backend/requirements.txt

@@ -8,4 +8,5 @@ pydantic_settings
 aiomysql
 greenlet
 werkzeug
-python-multipart
+python-multipart
+httpx

backend/service/log_service.py

@@ -30,7 +30,7 @@ def get_last_n_lines(file_path: str, n: int = 100) -> str:
             # Move to the end of the file
             f.seek(0, os.SEEK_END)
             block_size = 1024
-            lines_found = deque()
+            lines_found: deque[str] = deque()
 
             while f.tell() > 0 and len(lines_found) <= n:
                 # Calculate the position and size of the next block to read
@@ -59,7 +59,7 @@ def get_last_n_lines(file_path: str, n: int = 100) -> str:
 
             return "\n".join(list(lines_found)[-n:])
     except Exception as e:
-        logger.error(f"Failed to read log file:- {str(e)}")
+        logger.error(f"Failed to read log file: {str(e)}")
         return ""
 
 
@@ -110,12 +110,12 @@ async def get_service_log_svc(service_name: str, offset: int, tail: int):
     log_file_path = os.path.join(LOG_DIR, f"{service_name}.log")
 
     if not os.path.exists(log_file_path):
-        logger.warning("Log file not found.")
+        logger.warning(f"Log file not found: {log_file_path}")
         return JSONResponse(
             status_code=404,
             content={
                 "status": "error",
-                "error": f"Log file for service '{service_name}' not found at {log_file_path}",
+                "error": f"Log file for service '{service_name}' not found",
             },
         )
     try: