adai/MemoryBear
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(prompt-optimizer): add prompt optimization APIs and database tables

Browse Source 
- Added API endpoints for prompt optimization:
  * POST /prompt/sessions: Create a new prompt optimization session
  * GET /prompt/sessions/{session_id}: Retrieve session message history
  * POST /prompt/sessions/{session_id}/messages: Send message and get optimized prompt
  * PUT /prompt/model: Create or update system prompt model configuration

- Added database models for prompt optimization:
  * prompt_opt_session: Stores session metadata
  * prompt_opt_session_history: Stores session message history
  * prompt_opt_message: Stores user and assistant messages
  * prompt_opt_model_config: Stores system prompt model configurations

- Updated service layer to handle message creation, prompt optimization, and variable parsing
- Added corresponding Pydantic schemas for request and response validation
This commit is contained in:
mengyonghao
2025-12-17 15:56:33 +08:00
parent 7386ea32f1
commit 64d9dde209
8 changed files with 963 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
api/app/controllers/__init__.py
View File

@@ -28,6 +28,7 @@ from . import (
public_share_controller,
multi_agent_controller,
workflow_controller,
prompt_optimizer_controller
)
# 创建管理端 API 路由器
@@ -58,5 +59,6 @@ manager_router.include_router(public_share_controller.router) # 公开路由（
manager_router.include_router(memory_dashboard_controller.router)
manager_router.include_router(multi_agent_controller.router)
manager_router.include_router(workflow_controller.router)
manager_router.include_router(prompt_optimizer_controller.router)
__all__ = ["manager_router"]

170
api/app/controllers/prompt_optimizer_controller.py Normal file
View File

@@ -0,0 +1,170 @@
import uuid
from fastapi import APIRouter, Depends, Path
from sqlalchemy.orm import Session
from app.core.logging_config import get_api_logger
from app.core.response_utils import success
from app.dependencies import get_current_user, get_db
from app.models.prompt_optimizer_model import RoleType
from app.schemas.prompt_optimizer_schema import PromptOptMessage, PromptOptModelSet, CreateSessionResponse, \
OptimizePromptResponse, SessionHistoryResponse, SessionMessage
from app.schemas.response_schema import ApiResponse
from app.services.prompt_optimizer_service import PromptOptimizerService
router = APIRouter(prefix="/prompt", tags=["Prompts-Optimization"])
logger = get_api_logger()
@router.post(
"/sessions",
summary="Create a new prompt optimization session",
response_model=ApiResponse
)
def create_prompt_session(
db: Session = Depends(get_db),
current_user=Depends(get_current_user),
):
"""
Create a new prompt optimization session for the current user.
Returns:
ApiResponse: Contains the newly generated session ID.
"""
service = PromptOptimizerService(db)
# create new session
session = service.create_session(current_user.tenant_id, current_user.id)
result_schema = CreateSessionResponse.model_validate(session)
return success(data=result_schema)
@router.get(
"/sessions/{session_id}",
summary="获取 prompt 优化历史对话",
response_model=ApiResponse
)
def get_prompt_session(
session_id: uuid.UUID = Path(..., description="Session ID"),
db: Session = Depends(get_db),
current_user=Depends(get_current_user),
):
"""
Retrieve all messages from a specified prompt optimization session.
Args:
session_id (UUID): The ID of the session to retrieve
db (Session): Database session
current_user: Current logged-in user
Returns:
ApiResponse: Contains the session ID and the list of messages.
"""
service = PromptOptimizerService(db)
history = service.get_session_message_history(
session_id=session_id,
user_id=current_user.id
)
messages = [
SessionMessage(role=role, content=content)
for role, content in history
]
result = SessionHistoryResponse(
session_id=session_id,
messages=messages
)
return success(data=result)
@router.post(
"/sessions/{session_id}/messages",
summary="Get prompt optimization",
response_model=ApiResponse
)
async def get_prompt_opt(
session_id: uuid.UUID = Path(..., description="Session ID"),
data: PromptOptMessage = ...,
db: Session = Depends(get_db),
current_user=Depends(get_current_user),
):
"""
Send a user message in the specified session and return the optimized prompt
along with its description and variables.
Args:
session_id (UUID): The session ID
data (PromptOptMessage): Contains the user message, model ID, and current prompt
db (Session): Database session
current_user: Current user information
Returns:
ApiResponse: Contains the optimized prompt, description, and a list of variables.
"""
service = PromptOptimizerService(db)
service.create_message(
tenant_id=current_user.tenant_id,
session_id=session_id,
user_id=current_user.id,
role=RoleType.USER,
content=data.message
)
opt_result = await service.optimize_prompt(
tenant_id=current_user.tenant_id,
model_id=data.model_id,
session_id=session_id,
user_id=current_user.id,
current_prompt=data.current_prompt,
message=data.message
)
service.create_message(
tenant_id=current_user.tenant_id,
session_id=session_id,
user_id=current_user.id,
role=RoleType.ASSISTANT,
content=opt_result.desc
)
variables = service.parser_prompt_variables(opt_result.prompt)
result = {
"prompt": opt_result.prompt,
"desc": opt_result.desc,
"variables": variables
}
result_schema = OptimizePromptResponse.model_validate(result)
return success(data=result_schema)
@router.put(
"/model",
summary="Create or update prompt model config",
response_model=ApiResponse
)
def set_system_prompt(
data: PromptOptModelSet = ...,
db: Session = Depends(get_db),
current_user=Depends(get_current_user),
):
"""
Create or update a system prompt model configuration for the tenant.
Args:
data (PromptOptModelSet): Model configuration data including model ID,
system prompt, and optional configuration ID
db (Session): Database session
current_user: Current user information
Returns:
UUID: The ID of the created or updated model configuration.
"""
if data.id is None:
data.id = uuid.uuid4()
model_config = PromptOptimizerService(db).create_update_model_config(
current_user.tenant_id,
data.id, data.model_id,
data.system_prompt
)
return success(data=model_config.id)

6
api/app/models/__init__.py
View File

@@ -20,6 +20,7 @@ from .data_config_model import DataConfig
from .multi_agent_model import MultiAgentConfig, AgentInvocation
from .workflow_model import WorkflowConfig, WorkflowExecution, WorkflowNodeExecution
from .retrieval_info import RetrievalInfo
from .prompt_optimizer_model import PromptOptimizerModelConfig, PromptOptimizerSession, PromptOptimizerSessionHistory
__all__ = [
"Tenants",
@@ -54,5 +55,8 @@ __all__ = [
"WorkflowConfig",
"WorkflowExecution",
"WorkflowNodeExecution",
"RetrievalInfo"
"RetrievalInfo",
"PromptOptimizerModelConfig",
"PromptOptimizerSession",
"PromptOptimizerSessionHistory"
]

19
api/app/models/models_model.py
View File

@@ -15,6 +15,25 @@ class ModelType(StrEnum):
EMBEDDING = "embedding"
RERANK = "rerank"
@classmethod
def from_str(cls, value: str) -> "ModelType":
"""
Get a ModelType enum instance from a string value.
Args:
value (str): The string representation of the model type.
Returns:
ModelType: The corresponding ModelType enum object.
Raises:
ValueError: If the given value does not match any ModelType.
"""
try:
return cls(value)
except ValueError:
raise ValueError(f"Invalid ModelType: {value}")
class ModelProvider(StrEnum):
"""模型提供商枚举"""

176
api/app/models/prompt_optimizer_model.py Normal file
View File

@@ -0,0 +1,176 @@
import datetime
import uuid
from enum import StrEnum
from sqlalchemy import Column, ForeignKey, Text, DateTime, String, Index
from sqlalchemy.dialects.postgresql import UUID
from app.db import Base
class RoleType(StrEnum):
"""
Enumeration of message roles used in prompt optimization conversations.
This enum standardizes the role identifiers for messages stored in the
prompt optimization session history, ensuring consistency across
system-generated messages, user inputs, and assistant responses.
Attributes:
SYSTEM (str): Represents system-level instructions or prompts that
define the behavior or constraints of the assistant.
USER (str): Represents messages originating from the end user.
ASSISTANT (str): Represents messages generated by the AI assistant.
"""
SYSTEM = "system"
USER = "user"
ASSISTANT = "assistant"
class PromptOptimizerModelConfig(Base):
"""
Prompt Optimization Model Configuration.
This table stores system-level prompt configurations for each tenant.
The configuration defines the base system prompt used during prompt
optimization sessions and serves as a foundational instruction set
for the optimization process.
Each tenant may have one or more model configurations depending on
business requirements.
Table Name:
prompt_model_config
Columns:
id (UUID):
Primary key. Unique identifier for the prompt model configuration.
tenant_id (UUID):
Foreign key referencing `tenants.id`.
Identifies the tenant that owns this configuration.
system_prompt (Text):
The system-level prompt used to guide prompt optimization logic.
created_at (DateTime):
Timestamp indicating when the configuration was created.
updated_at (DateTime):
Timestamp indicating the last update time of the configuration.
Usage:
- Loaded when initializing a prompt optimization session
- Acts as the root system instruction for all subsequent prompts
"""
__tablename__ = "prompt_model_config"
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4, index=True)
tenant_id = Column(UUID(as_uuid=True), ForeignKey("tenants.id"), nullable=False, comment="Tenant ID")
# model_id = Column(UUID(as_uuid=True), nullable=False, comment="Model ID")
system_prompt = Column(Text, nullable=False, comment="System Prompt")
created_at = Column(DateTime, default=datetime.datetime.now, comment="Creation Time")
updated_at = Column(DateTime, default=datetime.datetime.now, onupdate=datetime.datetime.now, comment="Update Time")
class PromptOptimizerSession(Base):
"""
Prompt Optimization Session Registry.
This table records high-level metadata for prompt optimization sessions.
Each record represents a single logical session initiated by a user
under a specific tenant.
The session acts as a container for multiple conversation messages
stored in the session history table.
Table Name:
prompt_opt_session_list
Columns:
id (UUID):
Primary key. Internal unique identifier for the session record.
tenant_id (UUID):
Foreign key referencing `tenants.id`.
Identifies the tenant under which the session is created.
session_id (UUID):
Public-facing session identifier used to group conversation history.
user_id (UUID):
Foreign key referencing `users.id`.
Identifies the user who initiated the session.
created_at (DateTime):
Timestamp indicating when the session was created.
Design Notes:
- This table intentionally does not store message content
- Message-level data is stored in `prompt_opt_session_history`
- Enables efficient session listing and pagination
"""
__tablename__ = "prompt_opt_session_list"
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4, index=True)
tenant_id = Column(UUID(as_uuid=True), ForeignKey("tenants.id"), nullable=False, comment="Tenant ID")
# app_id = Column(UUID(as_uuid=True), ForeignKey("apps.id"), nullable=False, comment="Application ID")
session_id = Column(UUID(as_uuid=True), nullable=False, comment="Session ID")
user_id = Column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=False, comment="User ID")
created_at = Column(DateTime, default=datetime.datetime.now, comment="Creation Time", index=True)
class PromptOptimizerSessionHistory(Base):
"""
Prompt Optimization Session Message History.
This table stores the complete conversational history of a prompt
optimization session, including system prompts, user inputs, and
assistant responses.
Each record represents a single message within a session, preserving
the chronological order of interactions.
Table Name:
prompt_opt_session_history
Columns:
id (UUID):
Primary key. Unique identifier for the message record.
tenant_id (UUID):
Foreign key referencing `tenants.id`.
Identifies the tenant under which the session operates.
session_id (UUID):
Logical session identifier linking messages to a session.
user_id (UUID):
Foreign key referencing `users.id`.
Identifies the user associated with the session.
message_role (Text):
Role of the message sender (e.g., system, user, assistant).
message_content (Text):
Raw message content generated or provided during the session.
prompt (Text):
The prompt snapshot used at the time of message generation.
created_at (DateTime):
Timestamp indicating when the message was created.
Design Notes:
- Supports full conversation replay and audit
- Enables prompt evolution tracking over time
- Indexed by creation time for efficient chronological queries
"""
__tablename__ = "prompt_opt_session_history"
__table_args__ = (
Index(
"ix_prompt_opt_session_history_session_user_created",
"session_id",
"user_id",
"created_at"
),
)
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4, index=True)
tenant_id = Column(UUID(as_uuid=True), ForeignKey("tenants.id"), nullable=False, comment="Tenant ID")
# app_id = Column(UUID(as_uuid=True), ForeignKey("apps.id"), nullable=False, comment="Application ID")
session_id = Column(UUID(as_uuid=True), nullable=False, comment="Session ID")
user_id = Column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=False, comment="User ID")
role = Column(String, nullable=False, comment="Message Role")
content = Column(Text, nullable=False, comment="Message Content")
# prompt = Column(Text, nullable=False, comment="Prompt")
created_at = Column(DateTime, default=datetime.datetime.now, comment="Creation Time", index=True)

210
api/app/repositories/prompt_optimizer_repository.py Normal file
View File

@@ -0,0 +1,210 @@
import uuid
from typing import Optional
from sqlalchemy.orm import Session
from app.core.logging_config import get_db_logger
from app.models.prompt_optimizer_model import (
PromptOptimizerModelConfig,
PromptOptimizerSession, PromptOptimizerSessionHistory, RoleType
)
db_logger = get_db_logger()
class PromptOptimizerModelConfigRepository:
"""Repository for managing prompt optimizer model configurations."""
def __init__(self, db: Session):
self.db = db
def get_by_tenant_id(self, tenant_id: uuid.UUID) -> Optional[PromptOptimizerModelConfig]:
"""
Retrieve the prompt optimizer model configuration for a specific tenant.
Args:
tenant_id (uuid.UUID): The unique identifier of the tenant.
Returns:
Optional[PromptOptimizerModelConfig]: The model configuration if found, else None.
"""
db_logger.debug(f"Get prompt optimization model configuration: tenant_id={tenant_id}")
try:
config = self.db.query(PromptOptimizerModelConfig).filter(
PromptOptimizerModelConfig.tenant_id == tenant_id,
# PromptOptimizerModelConfig.model_id == model_id
).first()
if config:
db_logger.debug(f"Prompt optimization model configuration found: (ID: {config.id})")
else:
db_logger.debug(f"Prompt optimization model configuration not found: tenant_id={tenant_id}")
return config
except Exception as e:
db_logger.error(
f"Error retrieving prompt optimization model configuration: tenant_id={tenant_id} - {str(e)}")
raise
def get_by_config_id(self, tenant_id: uuid.UUID, config_id: uuid.UUID) -> Optional[PromptOptimizerModelConfig]:
"""
Retrieve a specific prompt optimizer model configuration by config ID and tenant ID.
Args:
tenant_id (uuid.UUID): The unique identifier of the tenant.
config_id (uuid.UUID): The unique identifier of the model configuration.
Returns:
Optional[PromptOptimizerModelConfig]: The model configuration if found, else None.
"""
db_logger.debug(f"Get prompt optimization model configuration: config_id={config_id}, tenant_id={tenant_id}")
try:
model = self.db.query(PromptOptimizerModelConfig).filter(
PromptOptimizerModelConfig.tenant_id == tenant_id,
PromptOptimizerModelConfig.id == config_id
).first()
if model:
db_logger.debug(f"Prompt optimization model configuration found: (ID: {model.id})")
else:
db_logger.debug(f"Prompt optimization model configuration not found: config_id={config_id}")
return model
except Exception as e:
db_logger.error(
f"Error retrieving prompt optimization model configuration: model_id={config_id} - {str(e)}")
raise
def create_or_update(
self,
config_id: uuid.UUID,
tenant_id: uuid.UUID,
system_prompt: str,
) -> Optional[PromptOptimizerModelConfig]:
"""
Create a new or update an existing prompt optimizer model configuration.
If a configuration with the given config_id exists, it updates its system_prompt.
Otherwise, it creates a new configuration record.
Args:
config_id (uuid.UUID): The unique identifier for the configuration.
tenant_id (uuid.UUID): The tenant's unique identifier.
system_prompt (str): The system prompt content for prompt optimization.
Returns:
Optional[PromptOptimizerModelConfig]: The created or updated model configuration.
"""
db_logger.debug(f"Create/Update prompt optimization model configuration: tenant_id={tenant_id}")
existing_config = self.get_by_config_id(tenant_id, config_id)
if existing_config:
existing_config.system_prompt = system_prompt
self.db.commit()
self.db.refresh(existing_config)
db_logger.debug(f"Prompt optimization model configuration update: ID:{config_id}")
return existing_config
else:
config = PromptOptimizerModelConfig(
id=config_id,
# model_id=model_id,
tenant_id=tenant_id,
system_prompt=system_prompt
)
self.db.add(config)
self.db.commit()
self.db.refresh(config)
db_logger.debug(f"Prompt optimization model configuration created: ID:{config.id}")
return config
class PromptOptimizerSessionRepository:
"""Repository for managing prompt optimization sessions and session history."""
def __init__(self, db: Session):
self.db = db
def create_session(
self,
tenant_id: uuid.UUID,
user_id: uuid.UUID
) -> PromptOptimizerSession:
"""
Create a new prompt optimization session for a user and app.
Args:
tenant_id (uuid.UUID): The unique identifier of the tenant.
user_id (uuid.UUID): The unique identifier of the user.
Returns:
PromptOptimizerSession: The newly created session object.
"""
db_logger.debug(f"Create prompt optimization session: tenant_id={tenant_id}, user_id={user_id}")
try:
session = PromptOptimizerSession(
tenant_id=tenant_id,
user_id=user_id,
session_id=uuid.uuid4(),
)
self.db.add(session)
self.db.commit()
self.db.refresh(session)
db_logger.debug(f"Prompt optimization session created: ID:{session.id}")
return session
except Exception as e:
db_logger.error(f"Error creating prompt optimization session: user_id={user_id} - {str(e)}")
raise
def get_session_history(
self,
session_id: uuid.UUID,
user_id: uuid.UUID
) -> list[type[PromptOptimizerSessionHistory]]:
"""
Retrieve all message history of a specific prompt optimization session.
Args:
session_id (uuid.UUID): The unique identifier of the session.
user_id (uuid.UUID): The unique identifier of the user.
Returns:
list[PromptOptimizerSessionHistory]: A list of session history records
ordered by creation time ascending.
"""
db_logger.debug(f"Get prompt optimization session history: "
f"user_id={user_id}, session_id={session_id}")
try:
history = self.db.query(PromptOptimizerSessionHistory).filter(
PromptOptimizerSessionHistory.session_id == session_id,
PromptOptimizerSessionHistory.user_id == user_id
).order_by(PromptOptimizerSessionHistory.created_at.asc()).all()
return history
except Exception as e:
db_logger.error(f"Error retrieving prompt optimization session history: session_id={session_id} - {str(e)}")
raise
def create_message(
self,
tenant_id: uuid.UUID,
session_id: uuid.UUID,
user_id: uuid.UUID,
role: RoleType,
content: str,
) -> PromptOptimizerSessionHistory:
"""
Create a new message in the session history.
This method is a placeholder for future implementation.
"""
try:
message = PromptOptimizerSessionHistory(
tenant_id=tenant_id,
session_id=session_id,
user_id=user_id,
role=role.value,
content=content,
)
self.db.add(message)
self.db.commit()
return message
except Exception as e:
db_logger.error(f"Error creating prompt optimization session history: session_id={session_id} - {str(e)}")
raise

99
api/app/schemas/prompt_optimizer_schema.py Normal file
View File

@@ -0,0 +1,99 @@
from pydantic import BaseModel, Field
from uuid import UUID
# =========================================
# API Request Schemas
# =========================================
class PromptOptMessage(BaseModel):
model_id: UUID = Field(
...,
description="Model ID"
)
message: str = Field(
...,
min_length=1,
description="User's input message"
)
current_prompt: str = Field(
default="",
description="currently optimized prompt"
)
class PromptOptModelSet(BaseModel):
id: UUID | None = Field(
default=None,
description="Configuration ID"
)
system_prompt: str = Field(
...,
description="System Prompt"
)
# =========================================
# Service Layer Results
# =========================================
class OptimizePromptResult(BaseModel):
prompt: str = Field(
...,
description="Optimized Prompt"
)
desc: str = Field(
...,
description="Description"
)
# =========================================
# API Response Schemas
# =========================================
class CreateSessionResponse(BaseModel):
model_config = {"from_attributes": True}
session_id: UUID = Field(
...,
description="Session ID"
)
class OptimizePromptResponse(BaseModel):
model_config = {"from_attributes": True}
prompt: str = Field(
...,
description="Optimized Prompt"
)
desc: str = Field(
...,
description="Description"
)
variables: list = Field(
...,
description="Variables"
)
class SessionMessage(BaseModel):
role: str = Field(
...,
description="Message role (user/assistant)"
)
content: str = Field(
...,
description="Message content"
)
class SessionHistoryResponse(BaseModel):
session_id: UUID = Field(
...,
description="Session ID"
)
messages: list[SessionMessage] = Field(
...,
description="List of messages in the session"
)

282
api/app/services/prompt_optimizer_service.py Normal file
View File

@@ -0,0 +1,282 @@
import json
import re
import uuid
from langchain_core.prompts import ChatPromptTemplate
from sqlalchemy.orm import Session
from app.core.error_codes import BizCode
from app.core.exceptions import BusinessException
from app.core.logging_config import get_business_logger
from app.core.models import RedBearModelConfig
from app.core.models.llm import RedBearLLM
from app.models import ModelConfig, ModelApiKey, ModelType, PromptOptimizerSessionHistory
from app.models.prompt_optimizer_model import (
PromptOptimizerModelConfig,
PromptOptimizerSession,
RoleType
)
from app.repositories.model_repository import ModelConfigRepository
from app.repositories.prompt_optimizer_repository import (
PromptOptimizerModelConfigRepository,
PromptOptimizerSessionRepository
)
from app.schemas.prompt_optimizer_schema import OptimizePromptResult
logger = get_business_logger()
class PromptOptimizerService:
def __init__(self, db: Session):
self.db = db
def get_model_config(
self,
tenant_id: uuid.UUID,
model_id: uuid.UUID
) -> tuple[PromptOptimizerModelConfig, ModelConfig]:
"""
Retrieve the prompt optimizer model configuration and model configuration.
This method retrieves the prompt optimizer model configuration associated
with the specified model ID and tenant. It also fetches the corresponding
model configuration.
Args:
tenant_id (uuid.UUID): The unique identifier of the tenant.
model_id (uuid.UUID): The unique identifier of the prompt optimization model.
Returns:
tuple[PromptOptimzerModelConfig, ModelConfig]:
A tuple containing the prompt optimizer model configuration
and the corresponding model configuration.
Raises:
BusinessException: If the prompt optimizer model configuration does not exist.
BusinessException: If the model configuration does not exist.
"""
prompt_config = PromptOptimizerModelConfigRepository(self.db).get_by_tenant_id(
tenant_id
)
if not prompt_config:
raise BusinessException("提示词模型配置不存在", BizCode.NOT_FOUND)
model = ModelConfigRepository.get_by_id(
self.db, model_id, tenant_id=tenant_id
)
if not model:
raise BusinessException("模型配置不存在", BizCode.MODEL_NOT_FOUND)
return prompt_config, model
def create_update_model_config(
self,
tenant_id: uuid.UUID,
config_id: uuid.UUID,
model_id: uuid.UUID,
system_prompt: str,
) -> PromptOptimizerModelConfig:
"""
Create or update a prompt optimizer model configuration.
This method creates a new prompt optimizer model configuration or updates
an existing one identified by the given configuration ID. The configuration
defines the system prompt used for prompt optimization.
Args:
tenant_id (uuid.UUID): The unique identifier of the tenant.
config_id (uuid.UUID): The unique identifier of the configuration to create or update.
model_id (uuid.UUID): The unique identifier of the model associated with this configuration.
system_prompt (str): The system prompt content used for prompt optimization.
Returns:
PromptOptimzerModelConfig: The created or updated prompt optimizer model configuration.
"""
prompt_config = PromptOptimizerModelConfigRepository(self.db).create_or_update(
config_id=config_id,
tenant_id=tenant_id,
system_prompt=system_prompt,
)
return prompt_config
def create_session(
self,
tenant_id: uuid.UUID,
user_id: uuid.UUID
) -> PromptOptimizerSession:
"""
Create a new prompt optimization session.
This method initializes a new prompt optimization session for the specified
tenant, application, and user, and persists it to the database.
Args:
tenant_id (uuid.UUID): The unique identifier of the tenant.
user_id (uuid.UUID): The unique identifier of the user.
Returns:
PromptOptimzerSession: The newly created prompt optimization session.
"""
session = PromptOptimizerSessionRepository(self.db).create_session(
tenant_id=tenant_id,
user_id=user_id
)
return session
def get_session_message_history(
self,
session_id: uuid.UUID,
user_id: uuid.UUID
) -> list[tuple[str, str]]:
"""
Retrieve the chronological message history for a prompt optimization session.
This method queries the database to fetch all messages associated with a
specific prompt optimization session for a given user. Messages are returned
in chronological order and typically include both user inputs and
model-generated responses.
Args:
session_id (uuid.UUID): The unique identifier of the prompt optimization session.
user_id (uuid.UUID): The unique identifier of the user associated with the session.
Returns:
list[tuple[str, str]]: A list of tuples representing messages. Each tuple contains:
- role (str): The role of the message sender, e.g., 'system', 'user', or 'assistant'.
- content (str): The content of the message.
"""
history = PromptOptimizerSessionRepository(self.db).get_session_history(
session_id=session_id,
user_id=user_id
)
messages = []
for message in history:
messages.append((message.role, message.content))
return messages
async def optimize_prompt(
self,
tenant_id: uuid.UUID,
model_id: uuid.UUID,
session_id: uuid.UUID,
user_id: uuid.UUID,
current_prompt: str,
message: str
) -> OptimizePromptResult:
"""
Optimize a prompt using a prompt optimizer LLM.
This method uses a configured prompt optimizer model to refine an existing
prompt based on the user's requirements. The optimized prompt is generated
according to predefined system rules, including Jinja2 variable syntax and
a strict JSON output format.
Args:
tenant_id (uuid.UUID): The unique identifier of the tenant.
model_id (uuid.UUID): The unique identifier of the prompt optimizer model.
session_id (uuid.UUID): The unique identifier of the prompt optimization session.
user_id (uuid.UUID): The unique identifier of the user associated with the session.
current_prompt (str): The original prompt to be optimized.
message (str): The user's requirements or modification instructions.
Returns:
dict: A dictionary containing the optimized prompt and the description
of changes, in the following format:
{
"prompt": "<optimized_prompt>",
"desc": "<change_description>"
}
Raises:
BusinessException: If the model response cannot be parsed as valid JSON
or does not conform to the expected output format.
"""
prompt_config, model_config = self.get_model_config(tenant_id, model_id)
session_history = self.get_session_message_history(session_id=session_id, user_id=user_id)
# Create LLM instance
api_config: ModelApiKey = model_config.api_keys[0]
llm = RedBearLLM(RedBearModelConfig(
model_name=api_config.model_name,
provider=api_config.provider,
api_key=api_config.api_key,
base_url=api_config.api_base
), type=ModelType.from_str(model_config.type))
# build message
messages = [
# init system_prompt
(RoleType.SYSTEM.value, prompt_config.system_prompt),
# base model limit
(RoleType.SYSTEM.value,
"Optimization Rules:\n"
"1. Fully adjust the prompt content according to the user's requirements.\n"
"2. When the user requests the insertion of variables, you must use Jinja2 syntax {{variable_name}} "
"(the variable name should be determined based on the user's requirement).\n"
"3. Keep the prompt logic clear and instructions explicit.\n"
"4. Ensure that the modified prompt can be directly used.\n\n"
"Output Requirements:\n"
"Provide the result in JSON format, containing exactly two fields:\n"
" - prompt: The modified prompt (string).\n"
" - desc: A response addressing the user's optimization request (string).")
]
messages.extend(session_history[:-1]) # last message is current message
user_message_template = ChatPromptTemplate.from_messages([
(RoleType.USER.value, "[current_prompt]\n{current_prompt}\n[user_require]\n{message}")
])
formatted_user_message = user_message_template.format(current_prompt=current_prompt, message=message)
messages.extend([(RoleType.USER.value, formatted_user_message)])
logger.info(f"Prompt optimization message: {messages}")
result = await llm.ainvoke(messages)
try:
data_dict = json.loads(result.content)
model_resp = OptimizePromptResult.model_validate(data_dict)
except Exception as e:
logger.error(f"Failed to parse model reponse to json - Error: {str(e)}", exc_info=True)
raise BusinessException("Failed to parse model response", BizCode.PARSER_NOT_SUPPORTED)
return model_resp
@staticmethod
def parser_prompt_variables(prompt: str):
try:
pattern = r'\{\{\s*([a-zA-Z_][a-zA-Z0-9_]*)\s*\}\}'
matches = re.findall(pattern, prompt)
variables = list(set(matches))
return variables
except Exception as e:
logger.error(f"Failed to parse prompt variables - Error: {str(e)}", exc_info=True)
raise BusinessException("Failed to parse prompt variables", BizCode.PARSER_NOT_SUPPORTED)
@staticmethod
def fill_prompt_variables(prompt: str, variables: dict[str, str]):
try:
pattern = r'\{\{\s*([a-zA-Z_][a-zA-Z0-9_]*)\s*\}\}'
def replace_var(match):
var_name = match.group(1)
return variables.get(var_name, match.group(0))
result = re.sub(pattern, replace_var, prompt)
return result
except Exception as e:
logger.error(f"Failed to fill prompt variables - Error: {str(e)}", exc_info=True)
raise BusinessException("Failed to fill prompt variables", BizCode.PARSER_NOT_SUPPORTED)
def create_message(
self,
tenant_id: uuid.UUID,
session_id: uuid.UUID,
user_id: uuid.UUID,
role: RoleType,
content: str
) -> PromptOptimizerSessionHistory:
"""Insert Message to Session History"""
message = PromptOptimizerSessionRepository(self.db).create_message(
tenant_id=tenant_id,
session_id=session_id,
user_id=user_id,
role=role,
content=content
)
return message