common package

Submodules

common.basics module

class common.basics.DesignGoal(text: str | None = None, images: List[str] | None = None, extra: Dict[str, Any] | None = None)

Bases: object

Represents a design goal, which can be defined by either text, images, or both.

A design goal encapsulates the user’s input, which can be in the form of a textual description, one or more images, or a combination of both. Images can be provided as paths to individual image files or as a path to a folder containing multiple images.

Parameters:

• text (Optional[str]) – A textual description of the design goal. Defaults to None.

• images (Optional[list[str]]) – A list of image file paths or a single folder path containing images. Defaults to None.

• extra (Optional[Dict[str, Any]]) – Additional information related to the design goal, such as original user input or decomposed part list, etc. Defaults to an empty dictionary.

Raises:

ValueError – If neither text nor images is provided.

text

The textual description of the design goal.

Type:

Optional[str]

images

A list of image file paths or a single folder path.

Type:

Optional[list[str]]

extra

Additional information related to the design goal, such as original user input or decomposed part list, etc.

Type:

Dict[str, Any]

classmethod from_dict(data: Dict[str, Any]) → DesignGoal

Create a DesignGoal object from a dictionary.

Parameters:

data (Dict[str, Any]) – A dictionary containing the design goal data.

Returns:

A DesignGoal object.

Return type:

DesignGoal

classmethod from_json(json_str: str) → DesignGoal

Create a DesignGoal object from a JSON string.

Parameters:

json_str (str) – A JSON string containing the design goal data.

Returns:

A DesignGoal object.

Return type:

DesignGoal

to_dict() → Dict[str, Any]

Convert the DesignGoal object to a dictionary.

Returns:

A dictionary representation of the DesignGoal object.

Return type:

Dict[str, Any]

to_json() → str

Convert the DesignGoal object to a JSON string.

Returns:

A JSON string representation of the DesignGoal object.

Return type:

str

class common.basics.PromptBuilder

Bases: object

A utility class for constructing prompts with text and images.

This class is designed to build a list of messages that can be used as input for models that accept multi-modal prompts (e.g., text and images). Messages can include system prompts, user prompts with text, and user prompts with images.

messages

A list of messages, where each message is a dictionary containing a role (“system” or “user”) and content (text or image data).

Type:

list

add_images(image_data: list[str] | str)

Add images to the messages.

Accepts a list of image paths or a single image path. Each image is converted to a base64-encoded string and added as a user message with image content.

Parameters:

image_data (list[str] | str) – A list of image paths or a single image path.

add_system_prompt(content)

Add a system prompt to the messages.

Parameters:

content (str) – The content of the system prompt.

add_text(content)

Add a user message with text content to the messages.

Parameters:

content (str) – The text content of the user message.

add_tools(tools: ToolRegistry, keep_existing: bool = False)

Add tools to the PromptBuilder.

This method allows adding a registry of tools to the prompt. If tools already exist, they can be merged with the new ones based on the keep_existing flag. If keep_existing is True, existing tools with conflicting names will be preserved. Otherwise, they will be overwritten.

Parameters:

tools (ToolRegistry) – The registry of tools to be used with the prompt.

add_user_prompt(content)

Add a user prompt with text content to the messages.

Parameters:

content (str) – The text content of the user prompt.

get_tools()

Get the tools set for the prompt.

Returns:

The registry of tools if any have been set.

Return type:

ToolRegistry

common.embed module

class common.embed.Embed(api_key: str, api_base_url: str, model_name: str, org_id: str, **model_kwargs)

Bases: ABC

embed(texts: List[str]) → List[List[float]]

Generate embeddings for a list of texts using the OpenAI API.

Parameters:

texts (List[str]) – A list of text strings to generate embeddings for.

Returns:

A list of embeddings, where each embedding is a list of floats.

Return type:

List[List[float]]

common.llm module

class common.llm.LanguageModel(api_key, api_base_url, model_name, org_id=None, **model_kwargs)

Bases: ABC

query(prompt, system_prompt=None, tools: ToolRegistry = None)

Query the language model with a prompt, optionally a system prompt, and tools for function calling.

Parameters:

• prompt (str) – The user prompt to send to the model.

• system_prompt (str, optional) – The system prompt to guide the model’s behavior.

• tools (ToolRegistry, optional) – A registry of tools for function calling.

Returns:

The model’s response, including any function call results.

Return type:

str

query_with_promptbuilder(pb: PromptBuilder) → str

Query the LanguageModel using a PromptBuilder object.

Parameters:

pb – The PromptBuilder object containing the prompt.

Returns:

Generated response from the model.

Return type:

str

common.rerank module

class common.rerank.Rerank(api_key: str, api_base_url: str = 'https://api.siliconflow.cn/v1', model_name: str = 'BAAI/bge-reranker-v2-m3', **model_kwargs)

Bases: ABC

rerank(query: str, documents: List[str], top_n: int = 4, return_documents: bool = False) → List[Dict]

Rerank a list of documents based on a query.

Parameters:

• query (str) – The query to rerank documents against.

• documents (List[str]) – List of documents to rerank.

• top_n (int, optional) – Number of top documents to return. Defaults to 4.

• return_documents (bool, optional) – Whether to return the full documents or just scores. Defaults to False.

Returns:

List of reranked documents or scores.

Return type:

List[Dict]

common.tools module

class common.tools.Tool(*, name: str, description: str, parameters: Dict[str, Any], callable: Callable)

Bases: BaseModel

Represents a tool (function) that can be called by the language model.

callable: Callable

description: str

model_config: ClassVar[ConfigDict] = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

name: str

parameters: Dict[str, Any]

class common.tools.ToolRegistry

Bases: object

A registry for managing tools (functions) and their metadata.

get_callable(function_name: str) → Callable

Get a callable function by its name.

Parameters:

function_name (str) – The name of the function.

Returns:

The function to call, or None if not found.

Return type:

Callable

get_tools_json() → List[Dict[str, Any]]

Get the JSON representation of all registered tools, following JSON Schema.

Returns:

A list of tools in JSON format, compliant with JSON Schema.

Return type:

List[Dict[str, Any]]

merge(other: ToolRegistry, keep_existing: bool = False)

Merge tools from another ToolRegistry into this one.

Parameters:

other (ToolRegistry) – The other ToolRegistry to merge into this one.

register(func: Callable, description: str | None = None)

Register a function as a tool.

Parameters:

• func (Callable) – The function to register.

• description (str, optional) – A description of the function. If not provided, the function’s docstring will be used.

common.utils module

class common.utils.ANSIStrippingFormatter(fmt=None, datefmt=None, style='%', validate=True, *, defaults=None)

Bases: Formatter

A logging formatter that removes ANSI escape codes from log messages.

This formatter is particularly useful when logging to files, where ANSI escape codes (used for terminal coloring) are not needed and can clutter the log output.

ansi_escape = re.compile('\\x1B(?:[@-Z\\\\-_]|\\[[0-?]*[ -/]*[@-~])')

format(record)

Format the log record, removing any ANSI escape codes from the message.

Parameters:

record (logging.LogRecord) – The log record to format.

Returns:

The formatted log message with ANSI escape codes removed.

Return type:

str

common.utils.colorstring(message: Any, color: str | None = 'green', bold: bool = False) → str

Returns a colored string using either ANSI escape codes or blessed terminal capabilities.

Parameters:

• message – The message to be colored. Can be of any type (e.g., str, int, float, bool).

• color – The color to apply. Supported colors: ‘black’, ‘red’, ‘green’, ‘yellow’, ‘blue’, ‘magenta’, ‘cyan’, ‘white’, and their bright variants: ‘bright_black’, ‘bright_red’, ‘bright_green’, ‘bright_yellow’, ‘bright_blue’, ‘bright_magenta’, ‘bright_cyan’, ‘bright_white’.

• bold – If True, applies bold styling to the text (only applicable when use_ansi=True).

Returns:

A string with the specified color and styling.

common.utils.cprint(message: Any, color: str | None = 'green', **kwargs) → None

Prints a colored string using blessed terminal capabilities.

Parameters:

• message – The message to be colored. Can be of any type (e.g., str, int, float, bool).

• color – The color to apply. Supported colors: ‘black’, ‘red’, ‘green’, ‘yellow’, ‘blue’, ‘magenta’, ‘cyan’, ‘white’.

• kwargs – Additional keyword arguments to pass to the print function (e.g., end, sep, file, flush).

common.utils.extract_section_markdown(text: str, heading: str) → str

Extracts content from a markdown text under the specified heading.

common.utils.find_files_with_extensions(directory_path: str, extensions: str | Iterable[str], return_all: bool = False) → str | List[str] | None

Find files with the specified extensions in the given directory. If return_all is False (default), returns the first matching file based on priority. If return_all is True, returns a list of all matching files, sorted by priority.

Parameters:

• directory_path (str) – Path to the directory to search.

• extensions (Union[str, List[str]]) – A single extension or a list of extensions.

• return_all (bool) – If True, return all matching files; otherwise, return the first match.

Returns:

A single file path, a list of file paths, or None if no files are found.

Return type:

Union[str, List[str], None]

common.utils.get_image_paths(path: str | List[str]) → List[str]

Get image file paths from a specified folder, a single image file, or a list of image paths.

Parameters: path (Union[str, List[str]]): The path to the folder, the single image file, or a list of image paths.

Returns: List[str]: A list of image file paths.

Raises: ValueError: If any path does not exist or is not a valid image file or folder of images.

common.utils.image_to_base64(image: Image | str, quality: int = 85, max_resolution: tuple = (448, 448), img_format: str = 'WEBP') → str

Convert the image to a base64 encoded string.

Parameters:

• image – PIL Image object or the path to the image file.

• quality – Compression quality (0-100) for WebP format. Higher values mean better quality but larger size.

• max_resolution – Optional maximum resolution (width, height) to fit the image within while preserving aspect ratio.

• img_format – Image format to use for encoding. Default is “WEBP”.

Returns:

Base64 encoded string of the image.

common.utils.load_config(config_path: str, config_name: str | None = None) → dict

Load a YAML configuration file or a specific configuration from a folder or file.

Parameters:

• config_path (str) – Path to the YAML configuration file or folder containing configuration files.

• config_name (Optional[str]) – Name of the target configuration file (if config_path is a folder) or the key within the YAML file (if config_path is a file). If omitted and config_path is a file, the entire file is loaded.

Returns:

Dictionary containing the configuration data.

Return type:

dict

Raises:

• FileNotFoundError – If the specified config_path does not exist.

• yaml.YAMLError – If the YAML file is malformed or cannot be parsed.

• ValueError – If the config_name is not found in the configuration.

common.utils.load_prompts(prompts_path: str, which_model: str) → dict

Load prompts from a YAML file and return prompts for a specific model.

Parameters:

• prompts_path (str) – Path to the YAML file containing prompts.

• which_model (str) – Key specifying which model’s prompts to load.

Returns:

Dictionary containing prompts for the specified model.

Return type:

dict

Raises:

KeyError – If the specified which_model key is not found in the YAML file.

common.utils.parse_design_goal(design_goal_input: str) → str

Parse the design goal input, which can be either a JSON file or plain text. If it’s a JSON file, extract the ‘text’ field.

Parameters:

design_goal_input (str) – Path to a JSON file or plain text.

Returns:

The design goal text.

Return type:

str

common.utils.parse_json_response(response: str) → Dict[str, Any]

Parse JSON response from VLM, handling potential errors

common.utils.setup_logging(log_level: str = 'INFO', log_file: str | None = None, log_format: str = '%(asctime)s - %(name)s - %(levelname)s - %(message)s') → None

Configure the logging system.

This function sets up logging with the following behavior: - Console logs retain ANSI escape codes for colored output. - File logs (if specified) have ANSI escape codes removed for cleaner output. - Mutes httpx INFO logs while keeping the global log_level at INFO.

Parameters:

• log_level (str) – The logging level (e.g., “INFO”, “DEBUG”). Defaults to “INFO”.

• log_file (Optional[str]) – Path to the log file. If None, no file logging is performed.

• log_format (str) – The format string for log messages. Defaults to “%(asctime)s - %(name)s - %(levelname)s - %(message)s”.

Returns:

None

common.vlm module

class common.vlm.VisionLanguageModel(api_key: str, api_base_url: str, model_name: str, org_id: str, **model_kwargs)

Bases: LanguageModel, ABC

query_with_image(prompt: str | None = None, images: bytes | List[bytes] | None = None, images_with_text: List[str | bytes] | None = None, system_prompt: str | None = None) → str

Query the VisionLanguageModel with mixed text and image data.

Parameters:

• prompt – Optional user prompt text.

• images – Optional image data (single or list of bytes).

• images_with_text – Optional list of mixed text (str) and image (bytes) data.

• system_prompt – Optional system prompt text.

Returns:

Generated response from the model.

Module contents