feat: phase 2: workspace-level prompt self-modification
This commit is contained in:
parent
25aaa064f6
commit
5ac7b85ef7
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
),
|
||||
}
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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 ''}"""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
Loading…
Reference in New Issue