Skip to content

Commit

Permalink
Simplifies the document agent
Browse files Browse the repository at this point in the history 
By removing unused, redundant and overly complex code. And removing
UNDEFINED states and all of the checks they required.
  • Loading branch information
@markwaddle
markwaddle committed Jan 25, 2025
1 parent d3a59ff commit e2faa0c
Show file tree
Hide file tree
Showing 13 changed files with 382 additions and 1,159 deletions.
87 changes: 9 additions & 78 deletions assistants/prospector-assistant/assistant/agents/document/config.py
View file
Original file line number Diff line number Diff line change
@@ -1,59 +1,9 @@
import json
from typing import Annotated, Any, Dict, List, Type
from typing import Annotated

from guided_conversation.utils.resources import ResourceConstraint, ResourceConstraintMode, ResourceConstraintUnit
from pydantic import BaseModel, Field, create_model
from pydantic import BaseModel, Field
from semantic_workbench_assistant.config import UISchema

from ... import helpers
from . import config_defaults as config_defaults

#
# region Helpers
#

# take a full json schema and return a pydantic model, including support for
# nested objects and typed arrays


def json_type_to_python_type(json_type: str) -> Type:
# Mapping JSON types to Python types
type_mapping = {"integer": int, "string": str, "number": float, "boolean": bool, "object": dict, "array": list}
return type_mapping.get(json_type, Any)


def create_pydantic_model_from_json_schema(schema: Dict[str, Any], model_name="DynamicModel") -> Type[BaseModel]:
# Nested function to parse properties from the schema
def parse_properties(properties: Dict[str, Any]) -> Dict[str, Any]:
fields = {}
for prop_name, prop_attrs in properties.items():
prop_type = prop_attrs.get("type")
description = prop_attrs.get("description", None)

if prop_type == "object":
nested_model = create_pydantic_model_from_json_schema(prop_attrs, model_name=prop_name.capitalize())
fields[prop_name] = (nested_model, Field(..., description=description))
elif prop_type == "array":
items = prop_attrs.get("items", {})
if items.get("type") == "object":
nested_model = create_pydantic_model_from_json_schema(items)
fields[prop_name] = (List[nested_model], Field(..., description=description))
else:
nested_type = json_type_to_python_type(items.get("type"))
fields[prop_name] = (List[nested_type], Field(..., description=description))
else:
python_type = json_type_to_python_type(prop_type)
fields[prop_name] = (python_type, Field(..., description=description))
return fields

properties = schema.get("properties", {})
fields = parse_properties(properties)
return create_model(model_name, **fields)


# endregion


#
# region Models
#
Expand All @@ -69,46 +19,31 @@ class ResourceConstraintConfigModel(ResourceConstraint):
' "maximum", the agents will try to pace the conversation to use at most the resource quantity.'
),
),
] = config_defaults.resource_constraint.mode
]

unit: Annotated[
ResourceConstraintUnit,
Field(
title="Resource Unit",
description="The unit for the resource constraint.",
),
] = config_defaults.resource_constraint.unit
]

quantity: Annotated[
float,
Field(
title="Resource Quantity",
description="The quantity for the resource constraint. If <=0, the resource constraint is disabled.",
),
] = config_defaults.resource_constraint.quantity
]


class GuidedConversationConfigModel(BaseModel):
enabled: Annotated[
bool,
Field(description=helpers.load_text_include("guided_conversation_agent_enabled.md")),
UISchema(enable_markdown_in_description=True),
] = False

artifact: Annotated[
str,
Field(
title="Artifact",
description="The artifact that the agent will manage.",
),
UISchema(widget="baseModelEditor"),
] = json.dumps(config_defaults.ArtifactModel.model_json_schema(), indent=2)

rules: Annotated[
list[str],
Field(title="Rules", description="Do's and don'ts that the agent should attempt to follow"),
UISchema(schema={"items": {"ui:widget": "textarea", "ui:options": {"rows": 2}}}),
] = config_defaults.rules
]

conversation_flow: Annotated[
str,
Expand All @@ -117,7 +52,7 @@ class GuidedConversationConfigModel(BaseModel):
description="A loose natural language description of the steps of the conversation",
),
UISchema(widget="textarea", schema={"ui:options": {"rows": 10}}, placeholder="[optional]"),
] = config_defaults.conversation_flow.strip()
]

context: Annotated[
str,
Expand All @@ -126,19 +61,15 @@ class GuidedConversationConfigModel(BaseModel):
description="General background context for the conversation.",
),
UISchema(widget="textarea", placeholder="[optional]"),
] = config_defaults.context.strip()
]

resource_constraint: Annotated[
ResourceConstraintConfigModel,
Field(
title="Resource Constraint",
),
UISchema(schema={"quantity": {"ui:widget": "updown"}}),
] = ResourceConstraintConfigModel()

def get_artifact_model(self) -> Type[BaseModel]:
schema = json.loads(self.artifact)
return create_pydantic_model_from_json_schema(schema)
]


# endregion
62 changes: 0 additions & 62 deletions assistants/prospector-assistant/assistant/agents/document/config_defaults.py
View file

This file was deleted.

138 changes: 11 additions & 127 deletions ...stants/prospector-assistant/assistant/agents/document/gc_draft_content_feedback_config.py
View file
Original file line number Diff line number Diff line change
@@ -1,12 +1,6 @@
import json
from typing import Annotated, Any, Dict, List, Type
from guided_conversation.utils.resources import ResourceConstraintMode, ResourceConstraintUnit
from pydantic import BaseModel, Field

from guided_conversation.utils.resources import ResourceConstraint, ResourceConstraintMode, ResourceConstraintUnit
from pydantic import BaseModel, Field, create_model
from semantic_workbench_assistant.config import UISchema

from ... import helpers
from . import config_defaults as config_defaults
from .config import GuidedConversationConfigModel, ResourceConstraintConfigModel
from .guided_conversation import GC_ConversationStatus, GC_UserDecision

Expand Down Expand Up @@ -89,123 +83,13 @@ class ArtifactModel(BaseModel):
they might want or answer questions about it. This may be the first time the user is asking for you
help (conversation_status is user_initiated), or the nth time (conversation_status is user_returned)."""


# Resource Constraints (optional) - This defines the constraints on the conversation such as time or turns.
# It can also help with pacing the conversation,
# For example, here we have set an exact time limit of 10 turns which the agent will try to fill.
resource_constraint = ResourceConstraint(
quantity=5,
unit=ResourceConstraintUnit.TURNS,
mode=ResourceConstraintMode.MAXIMUM,
config = GuidedConversationConfigModel(
rules=rules,
conversation_flow=conversation_flow.strip(),
context=context.strip(),
resource_constraint=ResourceConstraintConfigModel(
unit=ResourceConstraintUnit.TURNS,
mode=ResourceConstraintMode.MAXIMUM,
quantity=5,
),
)


#
# region Helpers
#

# take a full json schema and return a pydantic model, including support for
# nested objects and typed arrays


def json_type_to_python_type(json_type: str) -> Type:
# Mapping JSON types to Python types
type_mapping = {"integer": int, "string": str, "number": float, "boolean": bool, "object": dict, "array": list}
return type_mapping.get(json_type, Any)


def create_pydantic_model_from_json_schema(schema: Dict[str, Any], model_name="DynamicModel") -> Type[BaseModel]:
# Nested function to parse properties from the schema
def parse_properties(properties: Dict[str, Any]) -> Dict[str, Any]:
fields = {}
for prop_name, prop_attrs in properties.items():
prop_type = prop_attrs.get("type")
description = prop_attrs.get("description", None)

if prop_type == "object":
nested_model = create_pydantic_model_from_json_schema(prop_attrs, model_name=prop_name.capitalize())
fields[prop_name] = (nested_model, Field(..., description=description))
elif prop_type == "array":
items = prop_attrs.get("items", {})
if items.get("type") == "object":
nested_model = create_pydantic_model_from_json_schema(items)
fields[prop_name] = (List[nested_model], Field(..., description=description))
else:
nested_type = json_type_to_python_type(items.get("type"))
fields[prop_name] = (List[nested_type], Field(..., description=description))
else:
python_type = json_type_to_python_type(prop_type)
fields[prop_name] = (python_type, Field(..., description=description))
return fields

properties = schema.get("properties", {})
fields = parse_properties(properties)
return create_model(model_name, **fields)


# endregion


#
# region Models
#


class GCDraftContentFeedbackConfigModel(GuidedConversationConfigModel):
enabled: Annotated[
bool,
Field(description=helpers.load_text_include("guided_conversation_agent_enabled.md")),
UISchema(enable_markdown_in_description=True),
] = False

artifact: Annotated[
str,
Field(
title="Artifact",
description="The artifact that the agent will manage.",
),
UISchema(widget="baseModelEditor"),
] = json.dumps(ArtifactModel.model_json_schema(), indent=2)

rules: Annotated[
list[str],
Field(title="Rules", description="Do's and don'ts that the agent should attempt to follow"),
UISchema(schema={"items": {"ui:widget": "textarea", "ui:options": {"rows": 2}}}),
] = rules

conversation_flow: Annotated[
str,
Field(
title="Conversation Flow",
description="A loose natural language description of the steps of the conversation",
),
UISchema(widget="textarea", schema={"ui:options": {"rows": 10}}, placeholder="[optional]"),
] = conversation_flow.strip()

context: Annotated[
str,
Field(
title="Context",
description="General background context for the conversation.",
),
UISchema(widget="textarea", placeholder="[optional]"),
] = context.strip()

resource_constraint: Annotated[
ResourceConstraintConfigModel,
Field(
title="Resource Constraint",
),
UISchema(schema={"quantity": {"ui:widget": "updown"}}),
] = ResourceConstraintConfigModel(
unit=resource_constraint.unit,
quantity=resource_constraint.quantity,
mode=resource_constraint.mode,
)

def get_artifact_model(self) -> Type[BaseModel]:
schema = json.loads(self.artifact)
return create_pydantic_model_from_json_schema(schema)


# endregion
Loading

0 comments on commit e2faa0c

Please sign in to comment.