diff --git a/README.md b/README.md index b2bd8fe2..4c94abbc 100644 --- a/README.md +++ b/README.md @@ -180,10 +180,14 @@ kwargs = {} ## system prompt to use for the vision model custom_system_prompt = None - # to override # custom_system_prompt = "For the below pdf page, do something..something..." ## example +## system role to use for the vision model +custom_role = None +# to override +# custom_role = "user" ## example + ###################### Example for OpenAI ###################### model = "gpt-4o-mini" ## openai model os.environ["OPENAI_API_KEY"] = "" ## your-api-key @@ -236,7 +240,8 @@ async def main(): output_dir = "./output_test" ## directory to save the consolidated markdown file result = await zerox(file_path=file_path, model=model, output_dir=output_dir, - custom_system_prompt=custom_system_prompt,select_pages=select_pages, **kwargs) + custom_system_prompt=custom_system_prompt, select_pages=select_pages, + custom_role=custom_role, **kwargs) return result @@ -259,6 +264,7 @@ async def zerox( output_dir: Optional[str] = None, temp_dir: Optional[str] = None, custom_system_prompt: Optional[str] = None, + custom_role: Optional[str] = None, select_pages: Optional[Union[int, Iterable[int]]] = None, **kwargs ) -> ZeroxOutput: @@ -283,7 +289,9 @@ Parameters - **temp_dir** (str, optional): The directory to store temporary files, defaults to some named folder in system's temp directory. If already exists, the contents will be deleted before zerox uses it. - **custom_system_prompt** (str, optional): - The system prompt to use for the model, this overrides the default system prompt of zerox.Generally it is not required unless you want some specific behaviour. When set, it will raise a friendly warning. Defaults to None. + The system prompt to use for the model, this overrides the default system prompt of zerox. Generally it is not required unless you want some specific behaviour. When set, it will raise a friendly warning. Defaults to None. +- **custom_role** (str, optional): + The role assigned to the model can be customized, overriding the default system role. Typically, this isn't necessary unless you need to specify a particular role for a given LLM. If you choose to set it, a friendly warning will be displayed. By default, this option is set to None. - **select_pages** (Optional[Union[int, Iterable[int]]], optional): Pages to process, can be a single page number or an iterable of page numbers, Defaults to None - **kwargs** (dict, optional): diff --git a/py_zerox/pyzerox/constants/messages.py b/py_zerox/pyzerox/constants/messages.py index e3ca8262..03a6cdb8 100644 --- a/py_zerox/pyzerox/constants/messages.py +++ b/py_zerox/pyzerox/constants/messages.py @@ -19,6 +19,10 @@ class Messages: Custom system prompt was provided which overrides the default system prompt. We assume that you know what you are doing. """ + CUSTOM_SYSTEM_ROLE_WARNING = """ + Custom system role was provided which overrides the default system role. We assume that you know what you are doing.\ + """ + MAINTAIN_FORMAT_SELECTED_PAGES_WARNING = """ The maintain_format flag is set to True in conjunction with select_pages input given. This may result in unexpected behavior. """ diff --git a/py_zerox/pyzerox/constants/prompts.py b/py_zerox/pyzerox/constants/prompts.py index 1b00f614..a59680a3 100644 --- a/py_zerox/pyzerox/constants/prompts.py +++ b/py_zerox/pyzerox/constants/prompts.py @@ -5,4 +5,4 @@ class Prompts: Convert the following PDF page to markdown. Return only the markdown with no explanation text. Do not exclude any content from the page. - """ \ No newline at end of file + """ diff --git a/py_zerox/pyzerox/constants/roles.py b/py_zerox/pyzerox/constants/roles.py new file mode 100644 index 00000000..b60ca923 --- /dev/null +++ b/py_zerox/pyzerox/constants/roles.py @@ -0,0 +1,4 @@ +class Roles: + """Class for storing roles for the Zerox system.""" + + DEFAULT_SYSTEM_ROLE = "system" diff --git a/py_zerox/pyzerox/core/types.py b/py_zerox/pyzerox/core/types.py index ffe251d3..3ac57acc 100644 --- a/py_zerox/pyzerox/core/types.py +++ b/py_zerox/pyzerox/core/types.py @@ -16,6 +16,7 @@ class ZeroxArgs: output_dir: Optional[str] = None temp_dir: Optional[str] = None custom_system_prompt: Optional[str] = None + custom_role: Optional[str] = None select_pages: Optional[Union[int, Iterable[int]]] = None kwargs: Dict[str, Any] = field(default_factory=dict) diff --git a/py_zerox/pyzerox/core/zerox.py b/py_zerox/pyzerox/core/zerox.py index 6b7c8dc1..addd1a88 100644 --- a/py_zerox/pyzerox/core/zerox.py +++ b/py_zerox/pyzerox/core/zerox.py @@ -34,6 +34,7 @@ async def zerox( output_dir: Optional[str] = None, temp_dir: Optional[str] = None, custom_system_prompt: Optional[str] = None, + custom_role: Optional[str] = None, select_pages: Optional[Union[int, Iterable[int]]] = None, **kwargs ) -> ZeroxOutput: @@ -57,6 +58,8 @@ async def zerox( :type temp_dir: str, optional :param custom_system_prompt: The system prompt to use for the model, this overrides the default system prompt of zerox. Generally it is not required unless you want some specific behaviour. When set, it will raise a friendly warning, defaults to None :type custom_system_prompt: str, optional + :param custom_role: The role assigned to the model can be customized, overriding the default system role. Typically, this isn't necessary unless you need to specify a particular role for a given LLM. If you choose to set it, a friendly warning will be displayed. By default, this option is set to None. + :type custom_role: str, optional :param select_pages: Pages to process, can be a single page number or an iterable of page numbers, defaults to None :type select_pages: int or Iterable[int], optional @@ -82,6 +85,10 @@ async def zerox( if custom_system_prompt: vision_model.system_prompt = custom_system_prompt + # override the system role if a custom role is provided + if custom_role: + vision_model.system_role = custom_role + # Check if both maintain_format and select_pages are provided if maintain_format and select_pages is not None: warnings.warn(Messages.MAINTAIN_FORMAT_SELECTED_PAGES_WARNING) @@ -199,4 +206,4 @@ async def zerox( input_tokens=input_token_count, output_tokens=output_token_count, pages=formatted_pages, - ) \ No newline at end of file + ) diff --git a/py_zerox/pyzerox/models/modellitellm.py b/py_zerox/pyzerox/models/modellitellm.py index e7b2d204..504f961f 100644 --- a/py_zerox/pyzerox/models/modellitellm.py +++ b/py_zerox/pyzerox/models/modellitellm.py @@ -9,14 +9,17 @@ from ..errors import ModelAccessError, NotAVisionModel, MissingEnvironmentVariables from ..constants.messages import Messages from ..constants.prompts import Prompts +from ..constants.roles import Roles from ..processor.image import encode_image_to_base64 DEFAULT_SYSTEM_PROMPT = Prompts.DEFAULT_SYSTEM_PROMPT +DEFAULT_SYSTEM_ROLE = Roles.DEFAULT_SYSTEM_ROLE class litellmmodel(BaseModel): ## setting the default system prompt _system_prompt = DEFAULT_SYSTEM_PROMPT + _system_role = DEFAULT_SYSTEM_ROLE def __init__( self, @@ -41,6 +44,11 @@ def __init__( def system_prompt(self) -> str: '''Returns the system prompt for the model.''' return self._system_prompt + + @property + def system_role(self) -> str: + '''Returns the system role for the model.''' + return self._system_role @system_prompt.setter def system_prompt(self, prompt: str) -> None: @@ -49,6 +57,15 @@ def system_prompt(self, prompt: str) -> None: ''' self._system_prompt = prompt + @system_role.setter + def system_role(self, role: str) -> None: + ''' + Sets/overrides the system role for the model. + Will raise a friendly warning to notify the user. + ''' + warnings.warn(f"{Messages.CUSTOM_SYSTEM_ROLE_WARNING}. Default role for zerox is: {DEFAULT_SYSTEM_ROLE}") + self._system_role = role + ## custom method on top of BaseModel def validate_environment(self) -> None: """Validates the environment variables required for the model.""" @@ -123,7 +140,7 @@ async def _prepare_messages( # Default system message messages: List[Dict[str, Any]] = [ { - "role": "system", + "role": self._system_role, "content": self._system_prompt, }, ] @@ -133,7 +150,7 @@ async def _prepare_messages( if maintain_format and prior_page: messages.append( { - "role": "system", + "role": self._system_role, "content": f'Markdown must maintain consistent formatting with the following page: \n\n """{prior_page}"""', }, )