Make model generation prompts more flexible

src/draive/generation/model/call.py

@@ -11,12 +11,13 @@
 ]
 
 
-async def generate_model[Generated: Model](
+async def generate_model[Generated: Model]( # noqa: PLR0913
  generated: type[Generated],
  /,
  *,
  instruction: str,
  input: MultimodalContent, # noqa: A002
+ schema_variable: str | None = None,
  tools: Toolbox | None = None,
  examples: Iterable[tuple[MultimodalContent, Generated]] | None = None,
  **extra: Any,
@@ -26,6 +27,7 @@ async def generate_model[Generated: Model](
  generated,
  instruction=instruction,
  input=input,
+ schema_variable=schema_variable,
  tools=tools or model_generation.tools,
  examples=examples,
  **extra,

src/draive/generation/model/generator.py

@@ -18,6 +18,7 @@ async def __call__[Generated: Model]( # noqa: PLR0913
  *,
  instruction: str,
  input: MultimodalContent, # noqa: A002
+ schema_variable: str | None = None,
  tools: Toolbox | None = None,
  examples: Iterable[tuple[MultimodalContent, Generated]] | None = None,
  **extra: Any,

src/draive/generation/model/lmm.py

@@ -10,23 +10,33 @@
 ]
 
 
-async def lmm_generate_model[Generated: Model](
+async def lmm_generate_model[Generated: Model]( # noqa: PLR0913
  generated: type[Generated],
  /,
  *,
  instruction: str,
  input: MultimodalContent, # noqa: A002
+ schema_variable: str | None = None,
  tools: Toolbox | None = None,
  examples: Iterable[tuple[MultimodalContent, Generated]] | None = None,
  **extra: Any,
 ) -> Generated:
- system_message: LMMMessage = LMMMessage(
- role="system",
- content=INSTRUCTION.format(
- instruction=instruction,
- format=generated.specification(),
- ),
- )
+ system_message: LMMMessage
+ if variable := schema_variable:
+ system_message = LMMMessage(
+ role="system",
+ content=instruction.format(**{variable: generated.specification()}),
+ )
+
+ else:
+ system_message = LMMMessage(
+ role="system",
+ content=DEFAULT_INSTRUCTION.format(
+ instruction=instruction,
+ schema=generated.specification(),
+ ),
+ )
+
  input_message: LMMMessage = LMMMessage(
  role="user",
  content=input,
@@ -71,13 +81,13 @@ async def lmm_generate_model[Generated: Model](
  return generated.from_json(completion.content_string)
 
 
-INSTRUCTION: str = """\
+DEFAULT_INSTRUCTION: str = """\
 {instruction}
 
 IMPORTANT!
 The result have to conform to the following JSON Schema:
 ```
-{format}
+{schema}
 ```
 Provide ONLY a single, raw, valid JSON without any comments or formatting.
 """

src/draive/helpers/logs.py

@@ -41,19 +41,17 @@ def setup_logging(
  },
  },
  "loggers": {
- "": { # root logger
+ name: {
  "handlers": ["console"],
  "level": "DEBUG" if debug else "INFO",
  "propagate": False,
- },
- **{
- name: {
- "handlers": ["console"],
- "level": "DEBUG" if debug else "INFO",
- "propagate": False,
- }
- for name in loggers
- },
+ }
+ for name in loggers
+ },
+ "root": { # root logger
+ "handlers": ["console"],
+ "level": "DEBUG" if debug else "INFO",
+ "propagate": False,
  },
  },
  )