Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions changelog.d/regenerate-corpus-icon-tool.added.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
- **Manually (re)generate a corpus icon via an agent.** A new
`regenerate_corpus_icon` agent tool
(`opencontractserver/llms/tools/core_tools/corpus_branding.py`) re-runs the
corpus logo generator on demand, so a user can ask the corpus chat agent to
give a collection a fresh icon at any time — not just at creation. It reuses
the auto-branding primitives (`_build_logo_prompt` + `agenerate_logo_image`
with the deterministic PIL monogram fallback) through a new public helper
`aregenerate_corpus_logo` in `opencontractserver/corpuses/services/branding.py`
and persists through the creator-gated `CorpusService.update_icon`. An optional
`additional_instructions` argument lets the agent steer the look (e.g. "use
blue tones and a gavel motif"); the hint is sanitised and length-capped
(`CORPUS_LOGO_ADDITIONAL_INSTRUCTIONS_MAX_CHARS`) before being folded into the
image prompt. The tool is registered in `tool_registry.py` (alias
`generate_corpus_icon`), `requires_approval` + `requires_corpus` +
`requires_write_permission`, creator-only, and is wired into the interactive
corpus agent's authenticated toolset
(`opencontractserver/llms/agents/pydantic_ai_agents.py`) alongside
`update_corpus_description`. Unlike the create-time path it deliberately
overwrites an existing icon and ignores `auto_branding_enabled` (a manual
regeneration is an explicit request). Tests:
`opencontractserver/tests/test_corpus_icon_tool.py`.
7 changes: 7 additions & 0 deletions opencontractserver/constants/corpus_branding.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,13 @@
CORPUS_LOGO_REQUEST_TIMEOUT_SECONDS = 90.0
CORPUS_LOGO_CONNECT_TIMEOUT_SECONDS = 10.0

# Upper bound on the free-text styling hint an agent/user may pass to the manual
# ``regenerate_corpus_icon`` tool. The hint is sanitised
# (``sanitize_plaintext_for_prompt``) and capped at this length before being
# appended to the image prompt, so a long crafted value cannot dominate the
# prompt or break out of it.
CORPUS_LOGO_ADDITIONAL_INSTRUCTIONS_MAX_CHARS = 500

# ---------------------------------------------------------------------------
# PIL monogram fallback
# ---------------------------------------------------------------------------
Expand Down
132 changes: 124 additions & 8 deletions opencontractserver/corpuses/services/branding.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,8 @@
if TYPE_CHECKING:
from opencontractserver.corpuses.models import Corpus
from opencontractserver.llms.api import ToolType
from opencontractserver.shared.services.conventions import ServiceResult
from opencontractserver.users.models import User

logger = logging.getLogger(__name__)

Expand Down Expand Up @@ -176,6 +178,84 @@ def _save() -> str:
return await _save()


async def aregenerate_corpus_logo(
corpus: Corpus,
user: User,
*,
additional_instructions: str | None = None,
) -> ServiceResult[None]:
"""Generate a fresh logo and persist it to ``corpus.icon`` (creator-gated).

The manual counterpart to :func:`_generate_logo`: it builds the same
text-to-image prompt (optionally augmented with caller-supplied
``additional_instructions``), generates the image via
:func:`agenerate_logo_image` (OpenAI Images with the deterministic PIL
monogram fallback), and writes it through
:meth:`CorpusService.update_icon`.

Unlike the auto-branding path this **deliberately overwrites** any existing
icon and ignores ``auto_branding_enabled`` — a manual regeneration is an
explicit request, not the create-time best-effort default. Authorisation is
still enforced: ``update_icon`` is creator-only, so the call is routed
through the freshly re-loaded acting user and a non-creator receives a
failure result with no write performed.

The corpus and user rows are re-loaded inside the save's
``database_sync_to_async`` boundary (mirroring :func:`_generate_logo`) so
the ORM write runs against that thread's connection and honours any state
change — icon upload, hard delete — that landed during the slow image
generation.

Returns the :class:`ServiceResult` from ``update_icon`` (success carries
``None``; failure carries a human-readable reason).
"""
from channels.db import database_sync_to_async

from opencontractserver.utils.image_generation import agenerate_logo_image

prompt = _build_logo_prompt(corpus, additional_instructions)
image_bytes, ext = await agenerate_logo_image(
prompt=prompt,
fallback_text=corpus.title or "Corpus",
fallback_seed=str(corpus.pk),
)

corpus_pk = corpus.pk
user_pk = user.pk

@database_sync_to_async
def _save() -> ServiceResult[None]:
# Deferred imports keep sync ORM access inside the
# database_sync_to_async boundary and avoid a circular import
# (corpus_service -> branding).
from django.contrib.auth import get_user_model

from opencontractserver.corpuses.models import Corpus
from opencontractserver.corpuses.services.corpus_service import CorpusService
from opencontractserver.shared.services.conventions import ServiceResult

try:
fresh = Corpus.objects.get(pk=corpus_pk)
except Corpus.DoesNotExist:
return ServiceResult.failure(
"Corpus no longer exists; the regenerated icon was not saved."
)

acting_user = get_user_model().objects.filter(pk=user_pk).first()
if acting_user is None:
return ServiceResult.failure(
"Acting user no longer exists; the regenerated icon was not saved."
)

# ``update_icon`` is the authoritative creator-only gate, so this helper
# stays safe even if a future caller skips an up-front check.
return CorpusService.update_icon(
acting_user, fresh, image_bytes=image_bytes, extension=ext
)

return await _save()


# --------------------------------------------------------------------------- #
# Prompt builders
# --------------------------------------------------------------------------- #
Expand Down Expand Up @@ -255,16 +335,49 @@ def _build_branding_system_prompt(corpus: Corpus, tools: list[str]) -> str:
return "\n".join(parts)


def _build_logo_prompt(corpus: Corpus) -> str:
def sanitize_logo_instruction_hint(additional_instructions: str | None) -> str:
"""Return the sanitised styling hint, or ``""`` if it adds nothing.

Shared by :func:`_build_logo_prompt` (which only appends the hint when this
is non-empty) and the ``regenerate_corpus_icon`` tool (which reports
``additional_instructions_applied`` from it) so the "was a hint applied?"
answer is derived from the *sanitised* value in one place. A value that is
blank or consists only of stripped characters (e.g. quotes) collapses to
``""`` here, so it never falsely reports as applied.
"""
if not (additional_instructions and additional_instructions.strip()):
return ""
from opencontractserver.constants.corpus_branding import (
CORPUS_LOGO_ADDITIONAL_INSTRUCTIONS_MAX_CHARS,
)
from opencontractserver.utils.prompt_sanitization import (
sanitize_plaintext_for_prompt,
)

return sanitize_plaintext_for_prompt(
additional_instructions.strip(),
max_length=CORPUS_LOGO_ADDITIONAL_INSTRUCTIONS_MAX_CHARS,
)


def _build_logo_prompt(
corpus: Corpus, additional_instructions: str | None = None
) -> str:
"""Text-to-image prompt for the corpus logo.

SECURITY: the title/description are user-controlled and are interpolated
directly into the (quoted) image prompt. A text-to-image model has no
``<user_content>`` fence concept, so we instead neutralise the values with
``sanitize_plaintext_for_prompt`` — stripping quotes and collapsing
whitespace — so a crafted title cannot break out of the quotes and inject
its own directives (e.g. ``" . Instead, render the text: ...``). This
mirrors the prompt-hardening applied to the README agent's system prompt.
``additional_instructions`` is an optional free-text styling hint supplied
by the manual ``regenerate_corpus_icon`` agent tool (the auto-branding path
passes ``None``). When present it is folded into the prompt so a user can
steer the look (e.g. "use blue tones and a gavel motif").

SECURITY: the title/description AND the styling hint are user-controlled and
are interpolated directly into the (quoted) image prompt. A text-to-image
model has no ``<user_content>`` fence concept, so we instead neutralise the
values with ``sanitize_plaintext_for_prompt`` — stripping quotes and
collapsing whitespace — so a crafted value cannot break out of the quotes
and inject its own directives (e.g. ``" . Instead, render the text: ...``).
This mirrors the prompt-hardening applied to the README agent's system
prompt.
"""
from opencontractserver.utils.prompt_sanitization import (
sanitize_plaintext_for_prompt,
Expand All @@ -283,6 +396,9 @@ def _build_logo_prompt(corpus: Corpus) -> str:
)
if description:
prompt += f" The collection is about: {description}."
hint = sanitize_logo_instruction_hint(additional_instructions)
if hint:
prompt += f" Additional style guidance: {hint}."
prompt += (
" Flat design, simple geometric shapes, a single focal symbol, bold "
"solid colors, centered on a plain background, no text, no words, no "
Expand Down
13 changes: 13 additions & 0 deletions opencontractserver/llms/agents/pydantic_ai_agents.py
Original file line number Diff line number Diff line change
Expand Up @@ -3413,6 +3413,19 @@ class DocAnswer(BaseModel):
if config.user_id is not None:
effective_tools.append(update_corpus_desc_tool_wrapped)

# Let the creator regenerate the corpus icon on demand — the same
# generator the auto-branding flow runs at creation time. Built from
# the registry so corpus_id/user_id are injected and the approval
# gate + creator-only write are honoured. Sibling of
# update_corpus_description (the other creator-only corpus-row write).
effective_tools.extend(
_build_tools_from_registry(
["regenerate_corpus_icon"],
corpus_id=context.corpus.id,
user_id=config.user_id,
)
)

# Deep-research tools: let the chat agent spawn a long-running,
# autonomous research job over this corpus (``start_deep_research``)
# and report on its progress (``check_deep_research_status``) so a
Expand Down
3 changes: 3 additions & 0 deletions opencontractserver/llms/tools/core_tools/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
apropose_caml_citation_match,
aread_corpus_caml_article,
)
from .corpus_branding import aregenerate_corpus_icon # noqa: F401
from .descriptions import ( # noqa: F401
aget_corpus_description,
aget_document_description,
Expand Down Expand Up @@ -188,6 +189,8 @@
"aapply_caml_article_edit",
"apropose_caml_citation_match",
"aread_corpus_caml_article",
# Corpus icon (logo) regeneration
"aregenerate_corpus_icon",
# Document indexing
"IndexEntryItem",
"acreate_document_index",
Expand Down
138 changes: 138 additions & 0 deletions opencontractserver/llms/tools/core_tools/corpus_branding.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,138 @@
"""Agent tool for manually (re)generating a corpus's icon/logo.

The corpus auto-branding flow
(``opencontractserver/corpuses/services/branding.py``) generates a logo once,
at corpus-creation time, when no icon was uploaded. This tool exposes the
*same* generator to an agent so a user can trigger it on demand — e.g.
"regenerate the icon for this collection" or "give it a logo featuring a set of
scales" — at any point in a corpus conversation.

Permissioning mirrors the rest of the corpus-row write surface:

* The corpus must be visible to the calling user (IDOR-safe error otherwise).
* Writing the icon is **creator-only**, exactly like
:meth:`CorpusService.update_icon` / ``update_corpus_description`` — even a
collaborator with a guardian UPDATE grant cannot replace it. The creator gate
is checked up front (so we never spend an image-generation round-trip on a
request that cannot succeed) and is re-enforced by ``update_icon`` itself.

The tool is registered ``requires_approval=True`` so each regeneration surfaces
a confirmation prompt before it overwrites the existing icon.

``corpus_id`` and ``user_id`` are framework-injected (hidden from the LLM) by
``build_inject_params_for_context``; the only LLM-visible argument is the
optional ``additional_instructions`` styling hint.
"""

from __future__ import annotations

import logging
from typing import Any

from django.contrib.auth import get_user_model

from opencontractserver.corpuses.models import Corpus

from ._helpers import _db_sync_to_async

logger = logging.getLogger(__name__)

User = get_user_model()


def _authorize_icon_regeneration(
corpus_id: int, user_id: int | None
) -> tuple[Corpus, Any]:
"""Resolve and authorize ``(corpus, user)`` for an icon regeneration.

The user element of the return is typed ``Any``: ``User`` here is the
``get_user_model()`` runtime variable (not a type alias), so a quoted
annotation would trip mypy's ``valid-type`` check (mirrors
``extracts_and_analyzers._get_user_or_none``).

Raises ``PermissionError`` — a security exception the tool wrapper
propagates rather than swallowing — when the user is anonymous, missing,
cannot see the corpus, or is not the corpus creator. The "cannot access"
branch returns the same opaque framing the other corpus tools use so the
message cannot be used to enumerate corpora.
"""
if user_id is None:
raise PermissionError("regenerate_corpus_icon requires an authenticated user.")

user = User.objects.filter(pk=user_id).first()
if user is None:
raise PermissionError(f"User {user_id} not found.")

corpus = Corpus.objects.visible_to_user(user).filter(pk=corpus_id).first()
if corpus is None:
raise PermissionError(f"User {user_id} cannot access corpus {corpus_id}.")

# Creator-only: mirror ``CorpusService.update_icon``'s authoritative gate so
# we fail fast (before generating an image) on a request that cannot persist.
if corpus.creator_id != user.id:
raise PermissionError(
f"Only the corpus creator can regenerate the icon for corpus {corpus_id}."
)
return corpus, user


async def aregenerate_corpus_icon(
*,
corpus_id: int,
user_id: int | None = None,
additional_instructions: str | None = None,
) -> dict[str, Any]:
"""Generate a fresh icon/logo for the corpus and save it.

Re-runs the corpus logo generator (OpenAI Images with a deterministic PIL
monogram fallback) and writes the result to the corpus's ``icon``,
**replacing** any existing icon. Useful when the creator wants a new look,
or wants to steer the style via ``additional_instructions``.

Args:
corpus_id: Corpus whose icon to regenerate (injected from context).
user_id: User performing the regeneration (injected from context);
must be the corpus creator.
additional_instructions: Optional free-text styling hint folded into
the image prompt (e.g. "use blue tones and a gavel motif"). Only
affects AI generation; the deterministic monogram fallback (used
when image generation is disabled or unconfigured) ignores it.
Sanitised and length-capped before use.

Returns:
A small status dict describing the result.
"""
corpus, user = await _db_sync_to_async(_authorize_icon_regeneration)(
corpus_id, user_id
)

# Deferred import: the branding service imports from the corpuses package at
# module load, so importing it at the top of a tool module that is itself
# imported during tool-registry population risks a cycle.
from opencontractserver.corpuses.services.branding import (
aregenerate_corpus_logo,
sanitize_logo_instruction_hint,
)

result = await aregenerate_corpus_logo(
corpus, user, additional_instructions=additional_instructions
)
if not result.ok:
# The up-front gate already screens the creator-only case, so reaching
# here means the corpus/user state changed mid-flight (e.g. a hard
# delete during the slow image generation). Surface it as an operational
# error string so the agent can inform the user rather than crashing the
# turn.
raise ValueError(result.error or "Failed to save the regenerated corpus icon.")

return {
"corpus_id": corpus.id,
"status": "updated",
# Report from the *sanitised* hint (the same value the prompt builder
# appends) so a hint that collapses to empty (e.g. only quotes) is not
# falsely reported as applied.
"additional_instructions_applied": bool(
sanitize_logo_instruction_hint(additional_instructions)
),
"detail": "A new corpus icon was generated and saved.",
}
Loading