feat: phase 2: workspace-level prompt self-modification

This commit is contained in:
Benjamin McCormick 2026-01-29 13:37:28 -05:00
parent 25aaa064f6
commit 5ac7b85ef7
10 changed files with 374 additions and 4 deletions

View File

@ -68,6 +68,8 @@ from .workspace import (
get_all_workspaces,
get_or_create_workspace,
get_workspace,
get_workspace_agent_config,
set_workspace_agent_config,
update_workspace,
)
@ -140,6 +142,8 @@ __all__ = [
"delete_workspace",
"get_or_create_workspace",
"get_workspace",
"get_workspace_agent_config",
"set_workspace_agent_config",
"get_all_workspaces",
"update_workspace",
]

View File

@ -15,6 +15,8 @@ from src.utils.filter import apply_filter
from src.utils.types import GetOrCreateResult
from src.vector_store import get_external_vector_store
AGENT_CONFIG_KEY = "_agent_config"
logger = getLogger(__name__)
@ -451,3 +453,52 @@ async def delete_workspace(
messages_deleted=messages_count,
conclusions_deleted=conclusions_count,
)
async def get_workspace_agent_config(
db: AsyncSession,
workspace_name: str,
) -> schemas.WorkspaceAgentConfig:
"""
Get the agent configuration for a workspace.
Args:
db: Database session
workspace_name: Name of the workspace
Returns:
WorkspaceAgentConfig with the workspace's custom rules,
or defaults if not configured
"""
workspace = await get_workspace(db, workspace_name)
agent_config_data = workspace.h_metadata.get(AGENT_CONFIG_KEY, {})
return schemas.WorkspaceAgentConfig.model_validate(agent_config_data)
async def set_workspace_agent_config(
db: AsyncSession,
workspace_name: str,
config: schemas.WorkspaceAgentConfig,
) -> models.Workspace:
"""
Set the agent configuration for a workspace.
Args:
db: Database session
workspace_name: Name of the workspace
config: The agent configuration to set
Returns:
The updated workspace
"""
workspace = await get_workspace(db, workspace_name)
# Merge agent config into metadata
new_metadata = workspace.h_metadata.copy()
new_metadata[AGENT_CONFIG_KEY] = config.model_dump()
return await update_workspace(
db,
workspace_name,
schemas.WorkspaceUpdate(metadata=new_metadata),
)

View File

@ -70,6 +70,12 @@ async def process_representation_tasks_batch(
),
)
# Fetch workspace agent config for custom deriver rules
async with tracked_db("minimal_deriver.get_agent_config") as db:
agent_config = await crud.get_workspace_agent_config(
db, latest_message.workspace_name
)
# Skip if disabled
if message_level_configuration.reasoning.enabled is False:
return
@ -108,7 +114,11 @@ async def process_representation_tasks_batch(
)
# Build prompt
prompt = minimal_deriver_prompt(peer_id=observed, messages=formatted_messages)
prompt = minimal_deriver_prompt(
peer_id=observed,
messages=formatted_messages,
custom_rules=agent_config.deriver_rules,
)
context_prep_duration = (time.perf_counter() - overall_start) * 1000
accumulate_metric(

View File

@ -14,6 +14,7 @@ from src.utils.tokens import estimate_tokens
def minimal_deriver_prompt(
peer_id: str,
messages: str,
custom_rules: str = "",
) -> str:
"""
Generate minimal prompt for fast observation extraction.
@ -21,10 +22,19 @@ def minimal_deriver_prompt(
Args:
peer_id: The ID of the user being analyzed.
messages: All messages in the range (interleaving messages and new turns combined).
custom_rules: Optional workspace-specific rules to inject.
Returns:
Formatted prompt string for observation extraction.
"""
# Build optional custom rules section
custom_rules_section = ""
if custom_rules:
custom_rules_section = f"""
ADDITIONAL RULES (workspace-specific):
{custom_rules}
"""
return c(
f"""
Analyze messages from {peer_id} to extract **explicit atomic facts** about them.
@ -39,7 +49,7 @@ RULES:
- Observations should make sense on their own. Each observation will be used in the future to better understand {peer_id}.
- Extract ALL observations from {peer_id} messages, using others as context.
- Contextualize each observation sufficiently (e.g. "Ann is nervous about the job interview at the pharmacy" not just "Ann is nervous")
{custom_rules_section}
EXAMPLES:
- EXPLICIT: "I just had my 25th birthday last Saturday" "{peer_id} is 25 years old", "{peer_id}'s birthday is June 21st"
- EXPLICIT: "I took my dog for a walk in NYC" "{peer_id} has a dog", "{peer_id} lives in NYC"

View File

@ -61,6 +61,9 @@ async def agentic_chat(
db, workspace_name, observer=observer, observed=observed
)
# Get workspace agent config for custom dialectic rules
agent_config = await crud.get_workspace_agent_config(db, workspace_name)
# Create and run the dialectic agent
agent = DialecticAgent(
db=db,
@ -71,6 +74,7 @@ async def agentic_chat(
observer_peer_card=observer_peer_card,
observed_peer_card=observed_peer_card,
reasoning_level=reasoning_level,
custom_rules=agent_config.dialectic_rules,
)
response = await agent.answer(query)
@ -122,6 +126,9 @@ async def agentic_chat_stream(
db, workspace_name, observer=observer, observed=observed
)
# Get workspace agent config for custom dialectic rules
agent_config = await crud.get_workspace_agent_config(db, workspace_name)
# Create and run the dialectic agent
agent = DialecticAgent(
db=db,
@ -132,6 +139,7 @@ async def agentic_chat_stream(
observer_peer_card=observer_peer_card,
observed_peer_card=observed_peer_card,
reasoning_level=reasoning_level,
custom_rules=agent_config.dialectic_rules,
)
async for chunk in agent.answer_stream(query):

View File

@ -89,6 +89,7 @@ class DialecticAgent:
observed_peer_card: list[str] | None = None,
metric_key: str | None = None,
reasoning_level: ReasoningLevel = "low",
custom_rules: str = "",
):
"""
Initialize the dialectic agent.
@ -103,6 +104,7 @@ class DialecticAgent:
observed_peer_card: Biographical information about the observed peer
metric_key: Optional key for logging metrics (if provided, agent won't log separately)
reasoning_level: Level of reasoning to apply
custom_rules: Workspace-specific rules to inject into the prompt
"""
self.db: AsyncSession = db
self.workspace_name: str = workspace_name
@ -119,7 +121,11 @@ class DialecticAgent:
{
"role": "system",
"content": prompts.agent_system_prompt(
observer, observed, observer_peer_card, observed_peer_card
observer,
observed,
observer_peer_card,
observed_peer_card,
custom_rules,
),
}
]

View File

@ -8,6 +8,7 @@ def agent_system_prompt(
observed: str,
observer_peer_card: list[str] | None,
observed_peer_card: list[str] | None,
custom_rules: str = "",
) -> str:
"""
Generate the agent system prompt for the dialectic agent.
@ -17,6 +18,7 @@ def agent_system_prompt(
observed: The peer being queried about
observer_peer_card: Biographical information about the observer
observed_peer_card: Biographical information about the observed peer
custom_rules: Optional workspace-specific rules to inject
Returns:
Formatted system prompt string for the agent
@ -234,4 +236,8 @@ If after thorough searching you find NOTHING relevant:
After gathering context, reason through the information you found *before* stating your final answer. For comparison questions, explicitly compare the values. Only after you've verified your reasoning should you state your conclusion. Do NOT be pedantic, rather, be helpful and try to give the answer that the asker would expect -- they're the one who knows the most about themselves. Try to 'read their mind' -- understand the information they're really after and share it with them! Be **as specific as possible** given the information you have.
Do not explain your tool usage - just provide the synthesized answer.
"""
{f'''
## ADDITIONAL GUIDELINES (workspace-specific)
{custom_rules}
''' if custom_rules else ''}"""

View File

@ -92,6 +92,18 @@ async def update_workspace(
db: AsyncSession = db,
):
"""Update Workspace metadata and/or configuration."""
# Validate _agent_config if present in metadata
if workspace.metadata and "_agent_config" in workspace.metadata:
try:
schemas.WorkspaceAgentConfig.model_validate(
workspace.metadata["_agent_config"]
)
except Exception as e:
raise HTTPException(
status_code=400,
detail=f"Invalid _agent_config in metadata: {e}",
) from e
# ResourceNotFoundException will be caught by global handler if workspace not found
honcho_workspace = await crud.update_workspace(
db, workspace_name=workspace_id, workspace=workspace

View File

@ -27,6 +27,23 @@ class DreamType(str, Enum):
OMNI = "omni"
class WorkspaceAgentConfig(BaseModel):
"""
Configuration for per-workspace agent prompt customization.
Stored in workspace.metadata["_agent_config"].
"""
deriver_rules: str = Field(
default="",
description="Custom rules injected into the deriver prompt RULES section",
)
dialectic_rules: str = Field(
default="",
description="Custom rules injected into the dialectic prompt WORKFLOW section",
)
class ReconcilerType(str, Enum):
"""Types of reconciler tasks that can be performed."""

View File

@ -0,0 +1,246 @@
"""Tests for workspace agent config functionality (Phase 2 of Agentic FDE)."""
import pytest
from nanoid import generate as generate_nanoid
from sqlalchemy.ext.asyncio import AsyncSession
from src import crud, models
from src.deriver.prompts import minimal_deriver_prompt
from src.dialectic.prompts import agent_system_prompt
from src.schemas import WorkspaceAgentConfig
class TestWorkspaceAgentConfigSchema:
"""Test the WorkspaceAgentConfig Pydantic schema."""
def test_default_values(self):
"""Test that default values are empty strings."""
config = WorkspaceAgentConfig()
assert config.deriver_rules == ""
assert config.dialectic_rules == ""
def test_custom_values(self):
"""Test setting custom values."""
config = WorkspaceAgentConfig(
deriver_rules="Focus on technical facts only",
dialectic_rules="Be concise in responses",
)
assert config.deriver_rules == "Focus on technical facts only"
assert config.dialectic_rules == "Be concise in responses"
def test_model_dump(self):
"""Test serialization to dict."""
config = WorkspaceAgentConfig(
deriver_rules="Rule 1",
dialectic_rules="Rule 2",
)
data = config.model_dump()
assert data == {
"deriver_rules": "Rule 1",
"dialectic_rules": "Rule 2",
}
def test_model_validate(self):
"""Test deserialization from dict."""
data = {
"deriver_rules": "Custom rule",
"dialectic_rules": "Another rule",
}
config = WorkspaceAgentConfig.model_validate(data)
assert config.deriver_rules == "Custom rule"
assert config.dialectic_rules == "Another rule"
def test_model_validate_empty_dict(self):
"""Test deserialization from empty dict uses defaults."""
config = WorkspaceAgentConfig.model_validate({})
assert config.deriver_rules == ""
assert config.dialectic_rules == ""
class TestDeriverPromptCustomRules:
"""Test custom rules injection in deriver prompt."""
def test_empty_custom_rules_unchanged(self):
"""Test that empty custom_rules produces unchanged output."""
prompt_without = minimal_deriver_prompt(
peer_id="user123",
messages="Hello world",
)
prompt_with_empty = minimal_deriver_prompt(
peer_id="user123",
messages="Hello world",
custom_rules="",
)
assert prompt_without == prompt_with_empty
def test_custom_rules_injected(self):
"""Test that custom_rules are injected into prompt."""
custom_rules = "Focus on technical facts\nIgnore casual greetings"
prompt = minimal_deriver_prompt(
peer_id="user123",
messages="Hello world",
custom_rules=custom_rules,
)
assert "ADDITIONAL RULES (workspace-specific):" in prompt
assert "Focus on technical facts" in prompt
assert "Ignore casual greetings" in prompt
def test_custom_rules_position(self):
"""Test that custom rules appear after RULES section."""
custom_rules = "My custom rule"
prompt = minimal_deriver_prompt(
peer_id="user123",
messages="Hello world",
custom_rules=custom_rules,
)
# Find positions
rules_pos = prompt.find("RULES:")
custom_pos = prompt.find("ADDITIONAL RULES (workspace-specific):")
examples_pos = prompt.find("EXAMPLES:")
# Custom rules should be between RULES and EXAMPLES
assert rules_pos < custom_pos < examples_pos
class TestDialecticPromptCustomRules:
"""Test custom rules injection in dialectic prompt."""
def test_empty_custom_rules_unchanged(self):
"""Test that empty custom_rules produces no additional section."""
prompt_without = agent_system_prompt(
observer="agent",
observed="user123",
observer_peer_card=None,
observed_peer_card=None,
)
prompt_with_empty = agent_system_prompt(
observer="agent",
observed="user123",
observer_peer_card=None,
observed_peer_card=None,
custom_rules="",
)
# Both should not have the ADDITIONAL GUIDELINES section
assert "ADDITIONAL GUIDELINES (workspace-specific)" not in prompt_without
assert "ADDITIONAL GUIDELINES (workspace-specific)" not in prompt_with_empty
def test_custom_rules_injected(self):
"""Test that custom_rules are injected into prompt."""
custom_rules = "Always prioritize recent information\nBe skeptical of old data"
prompt = agent_system_prompt(
observer="agent",
observed="user123",
observer_peer_card=None,
observed_peer_card=None,
custom_rules=custom_rules,
)
assert "ADDITIONAL GUIDELINES (workspace-specific)" in prompt
assert "Always prioritize recent information" in prompt
assert "Be skeptical of old data" in prompt
def test_custom_rules_with_peer_cards(self):
"""Test that custom_rules work with peer cards enabled."""
custom_rules = "Custom rule here"
prompt = agent_system_prompt(
observer="agent",
observed="user123",
observer_peer_card=["Agent is helpful"],
observed_peer_card=["User likes coffee"],
custom_rules=custom_rules,
)
# Both peer cards and custom rules should be present
assert "Agent is helpful" in prompt
assert "User likes coffee" in prompt
assert "ADDITIONAL GUIDELINES (workspace-specific)" in prompt
assert "Custom rule here" in prompt
class TestWorkspaceAgentConfigCRUD:
"""Test CRUD operations for workspace agent config."""
@pytest.mark.asyncio
async def test_get_workspace_agent_config_default(
self,
db_session: AsyncSession,
sample_data: tuple[models.Workspace, models.Peer],
):
"""Test getting config for workspace without _agent_config returns defaults."""
workspace, _ = sample_data
config = await crud.get_workspace_agent_config(db_session, workspace.name)
assert config.deriver_rules == ""
assert config.dialectic_rules == ""
@pytest.mark.asyncio
async def test_set_and_get_workspace_agent_config(
self,
db_session: AsyncSession,
sample_data: tuple[models.Workspace, models.Peer],
):
"""Test setting and retrieving workspace agent config."""
workspace, _ = sample_data
# Set config
config = WorkspaceAgentConfig(
deriver_rules="Extract technical facts only",
dialectic_rules="Be concise",
)
await crud.set_workspace_agent_config(db_session, workspace.name, config)
# Get config
retrieved = await crud.get_workspace_agent_config(db_session, workspace.name)
assert retrieved.deriver_rules == "Extract technical facts only"
assert retrieved.dialectic_rules == "Be concise"
@pytest.mark.asyncio
async def test_set_workspace_agent_config_preserves_other_metadata(
self,
db_session: AsyncSession,
):
"""Test that setting agent config preserves existing metadata."""
# Create workspace with existing metadata
workspace_name = str(generate_nanoid())
workspace = models.Workspace(
name=workspace_name,
h_metadata={"custom_key": "custom_value", "another": 123},
)
db_session.add(workspace)
await db_session.flush()
# Set agent config
config = WorkspaceAgentConfig(deriver_rules="Test rule")
await crud.set_workspace_agent_config(db_session, workspace_name, config)
# Verify both old metadata and new config are present
await db_session.refresh(workspace)
assert workspace.h_metadata.get("custom_key") == "custom_value"
assert workspace.h_metadata.get("another") == 123
assert "_agent_config" in workspace.h_metadata
assert workspace.h_metadata["_agent_config"]["deriver_rules"] == "Test rule"
@pytest.mark.asyncio
async def test_update_workspace_agent_config(
self,
db_session: AsyncSession,
sample_data: tuple[models.Workspace, models.Peer],
):
"""Test updating workspace agent config."""
workspace, _ = sample_data
# Set initial config
config1 = WorkspaceAgentConfig(deriver_rules="Initial rule")
await crud.set_workspace_agent_config(db_session, workspace.name, config1)
# Update config
config2 = WorkspaceAgentConfig(
deriver_rules="Updated rule",
dialectic_rules="New dialectic rule",
)
await crud.set_workspace_agent_config(db_session, workspace.name, config2)
# Verify update
retrieved = await crud.get_workspace_agent_config(db_session, workspace.name)
assert retrieved.deriver_rules == "Updated rule"
assert retrieved.dialectic_rules == "New dialectic rule"