More Agent Utils (#36)

* More flexible * Reorder functions and imports * Changing the interface! * Chaining agent * Agent router * Export and bump version * Require agent's name * Better * Update * Update tests * Update examples
Rhobota · Nov 9, 2024 · b0a5c9f · b0a5c9f
1 parent ae82423
commit b0a5c9f
Show file tree

Hide file tree

Showing 15 changed files with 351 additions and 134 deletions.
diff --git a/examples/demo_committee.py b/examples/demo_committee.py
@@ -53,7 +53,7 @@ async def vote_on_jokes(joke_a: str, joke_b: str) -> Dict[str, int]:
         },
     ]
     tasks = [
-        model(noop_callback, [flat_messages(messages)])
+        model(noop_callback, [flat_messages('input', messages)])
         for model in COMMITTEE_MODELS
     ]
     outputs = await asyncio.gather(*tasks)

diff --git a/examples/demo_layered_agents.py b/examples/demo_layered_agents.py
@@ -1,6 +1,7 @@
 from lasagna import (
     known_models,
     build_simple_agent,
+    build_standard_message_extractor,
 )
 
 from lasagna.tui import (
@@ -35,17 +36,21 @@ async def main() -> None:
         tools = [
             evaluate_math_expression,
         ],
+        message_extractor = build_standard_message_extractor(
+            strip_tool_messages = True,
+            system_prompt_override = "You are a math assistant.",
+        ),
         doc = "Use this tool if the user asks a math question.",
-        system_prompt_override = "You are a math assistant.",
-        strip_old_tool_use_messages = True,
     )
     health_agent = known_models.BIND_ANTHROPIC_claude_35_sonnet()(
         build_simple_agent(
             name = 'health_agent',
             tools = [],
+            message_extractor = build_standard_message_extractor(
+                strip_tool_messages = True,
+                system_prompt_override = "You are a health coach who motivates through fear.",
+            ),
             doc = "Use this tool if the user asks a health question.",
-            system_prompt_override = "You are a health coach who motivates through fear.",
-            strip_old_tool_use_messages = True,
         ),
     )
     my_bound_agent = known_models.BIND_OPENAI_gpt_4o_mini()(

diff --git a/examples/demo_structured_output.py b/examples/demo_structured_output.py
@@ -41,14 +41,22 @@ class ExtractionModel(BaseModel):
 
 async def main() -> None:
     my_bound_agent = MODEL_BINDER(
-        build_extraction_agent(ExtractionModel),
+        build_extraction_agent(
+            name = 'extraction_agent',
+            extraction_type = ExtractionModel,
+        ),
     )
-    prev_runs = [flat_messages([
-        {
-            'role': 'human',
-            'text': PROMPT,
-        },
-    ])]
+    prev_runs = [
+        flat_messages(
+            'input',
+            [
+                {
+                    'role': 'human',
+                    'text': PROMPT,
+                },
+            ],
+        ),
+    ]
 
     agent_run = await my_bound_agent(tui_event_callback, prev_runs)
     print('DONE!')

diff --git a/examples/demo_vision.py b/examples/demo_vision.py
@@ -2,6 +2,7 @@
     known_models,
     build_simple_agent,
     flat_messages,
+    AgentRun,
     Message,
 )
 
@@ -35,7 +36,13 @@ async def main() -> None:
     bound_agent = MODEL_BINDER(
         build_simple_agent(name='agent'),
     )
-    await bound_agent(tui_event_callback, [flat_messages(messages)])
+    prev_runs: List[AgentRun] = [
+        flat_messages(
+            'input',
+            messages,
+        ),
+    ]
+    await bound_agent(tui_event_callback, prev_runs)
     print()
 
 

diff --git a/src/lasagna/__init__.py b/src/lasagna/__init__.py
@@ -4,23 +4,41 @@
 
 from .agent_util import (
     bind_model,
+    partial_bind_model,
     recursive_extract_messages,
+    extract_last_message,
     flat_messages,
+    override_system_prompt,
+    strip_tool_calls_and_results,
+    strip_all_but_last_human_message,
+    noop_callback,
+    MessageExtractor,
+    build_standard_message_extractor,
+    default_message_extractor,
     build_simple_agent,
     build_extraction_agent,
-    noop_callback,
-    extract_last_message,
+    build_agent_chainer,
+    build_agent_router,
 )
 
 __all__ = [
     'run',
     'bind_model',
+    'partial_bind_model',
     'recursive_extract_messages',
+    'extract_last_message',
     'flat_messages',
+    'override_system_prompt',
+    'strip_tool_calls_and_results',
+    'strip_all_but_last_human_message',
+    'noop_callback',
+    'MessageExtractor',
+    'build_standard_message_extractor',
+    'default_message_extractor',
     'build_simple_agent',
     'build_extraction_agent',
-    'noop_callback',
-    'extract_last_message',
+    'build_agent_chainer',
+    'build_agent_router',
 ]
 
-__version__ = "0.9.1"
+__version__ = "0.10.0"
diff --git a/src/lasagna/agent_runner.py b/src/lasagna/agent_runner.py
@@ -58,8 +58,6 @@ async def run(
         prev_runs,
     )
 
-    if 'agent' not in agent_run:
-        agent_run['agent'] = agent_name
     if 'provider' not in agent_run:
         agent_run['provider'] = provider_str
     if 'model' not in agent_run: