From 259361d279b7e372a1ec8ac8decf8c9ad529cf98 Mon Sep 17 00:00:00 2001 From: Alexander Myasoedov Date: Sat, 25 Jan 2025 12:50:02 +0200 Subject: [PATCH] feat(Add pydantic-ai agent): --- agentic_security/probe_actor/operator.py | 130 +++++++++++++++++++++-- agentic_security/probe_actor/refusal.py | 6 +- docs/external_module.md | 2 - docs/index.md | 2 - mkdocs.yml | 4 +- 5 files changed, 123 insertions(+), 21 deletions(-) diff --git a/agentic_security/probe_actor/operator.py b/agentic_security/probe_actor/operator.py index 5b8d31a..79b121d 100644 --- a/agentic_security/probe_actor/operator.py +++ b/agentic_security/probe_actor/operator.py @@ -1,30 +1,138 @@ +import asyncio +from typing import Any + +from pydantic_ai import Agent, RunContext + +# Define the OperatorToolBox class + + class OperatorToolBox: - def __init__(self, llm_spec, datasets): + def __init__(self, llm_spec: str, datasets: list[dict[str, Any]]): self.llm_spec = llm_spec self.datasets = datasets - def get_spec(self): + def get_spec(self) -> str: return self.llm_spec - def get_datasets(self): + def get_datasets(self) -> list[dict[str, Any]]: return self.datasets - def validate(self): + def validate(self) -> bool: # Validate the tool box - pass + return True - def stop(self): + def stop(self) -> None: # Stop the tool box pass - def run(self): + def run(self) -> None: # Run the tool box pass - def get_results(self): + def get_results(self) -> list[dict[str, Any]]: # Get the results - pass + return self.datasets - def get_failures(self): + def get_failures(self) -> list[str]: # Handle failure - pass + return [] + + def run_operation(self, operation: str) -> str: + # Run an operation + return f"Operation '{operation}' executed successfully." + + +# dataset_manager_agent.py + + +# Initialize OperatorToolBox +toolbox = OperatorToolBox( + llm_spec="GPT-4", datasets=["dataset1", "dataset2", "dataset3"] +) + +# Define the agent with OperatorToolBox as its dependency +dataset_manager_agent = Agent( + model="gpt-4", + deps_type=OperatorToolBox, + result_type=str, # The agent will return string results + system_prompt="You can validate the toolbox, run operations, and retrieve results or failures.", +) + + +@dataset_manager_agent.tool +async def validate_toolbox(ctx: RunContext[OperatorToolBox]) -> str: + """Validate the OperatorToolBox.""" + is_valid = ctx.deps.validate() + if is_valid: + return "ToolBox validation successful." + else: + return "ToolBox validation failed." + + +@dataset_manager_agent.tool +async def execute_operation(ctx: RunContext[OperatorToolBox], operation: str) -> str: + """Execute an operation on a dataset.""" + result = ctx.deps.run_operation(operation) + return result + + +@dataset_manager_agent.tool +async def retrieve_results(ctx: RunContext[OperatorToolBox]) -> str: + """Retrieve the results of operations.""" + results = ctx.deps.get_results() + if results: + formatted_results = "\n".join([f"{op}: {res}" for op, res in results.items()]) + return f"Operation Results:\n{formatted_results}" + else: + return "No operations have been executed yet." + + +@dataset_manager_agent.tool +async def retrieve_failures(ctx: RunContext[OperatorToolBox]) -> str: + """Retrieve the list of failures.""" + failures = ctx.deps.get_failures() + if failures: + formatted_failures = "\n".join(failures) + return f"Failures:\n{formatted_failures}" + else: + return "No failures recorded." + + +# Synchronous run example +def run_dataset_manager_agent_sync(): + prompts = [ + "Validate the toolbox.", + "Execute operation on 'dataset2'.", + "Execute operation on 'dataset4'.", # This should fail + "Retrieve the results.", + "Retrieve any failures.", + ] + + for prompt in prompts: + result = dataset_manager_agent.run_sync(prompt, deps=toolbox) + print(f"Prompt: {prompt}") + print(f"Response: {result.data}\n") + + +# Asynchronous run example +async def run_dataset_manager_agent_async(): + prompts = [ + "Validate the toolbox.", + "Execute operation on 'dataset2'.", + "Execute operation on 'dataset4'.", # This should fail + "Retrieve the results.", + "Retrieve any failures.", + ] + + for prompt in prompts: + result = await dataset_manager_agent.run(prompt, deps=toolbox) + print(f"Prompt: {prompt}") + print(f"Response: {result.data}\n") + + +if __name__ == "__main__": + # Run synchronous example + run_dataset_manager_agent_sync() + + # Run asynchronous example + asyncio.run(run_dataset_manager_agent_async()) diff --git a/agentic_security/probe_actor/refusal.py b/agentic_security/probe_actor/refusal.py index fe58c0d..a135423 100644 --- a/agentic_security/probe_actor/refusal.py +++ b/agentic_security/probe_actor/refusal.py @@ -1,5 +1,5 @@ -from typing import List, Dict from abc import ABC, abstractmethod + from agentic_security.refusal_classifier import RefusalClassifier classifier = RefusalClassifier() @@ -56,7 +56,7 @@ class RefusalClassifierPlugin(ABC): class DefaultRefusalClassifier(RefusalClassifierPlugin): """Default refusal classifier using a list of refusal phrases.""" - def __init__(self, refusal_phrases: List[str] = REFUSAL_MARKS): + def __init__(self, refusal_phrases: list[str] = REFUSAL_MARKS): self.refusal_phrases = refusal_phrases def is_refusal(self, response: str) -> bool: @@ -75,7 +75,7 @@ class RefusalClassifierManager: """Manager for refusal classifier plugins.""" def __init__(self): - self.plugins: Dict[str, RefusalClassifierPlugin] = {} + self.plugins: dict[str, RefusalClassifierPlugin] = {} def register_plugin(self, name: str, plugin: RefusalClassifierPlugin): """Register a refusal classifier plugin. diff --git a/docs/external_module.md b/docs/external_module.md index 4654b3c..3721b5c 100644 --- a/docs/external_module.md +++ b/docs/external_module.md @@ -4,7 +4,6 @@ The `Module` class interface provides a standardized way to create and use modul Here is an example of a module that implements the `ModuleProtocol` interface. This example shows how to create a module that processes prompts and sends results to a queue. - ```python from typing import List, Dict, Any, AsyncGenerator import asyncio @@ -42,4 +41,3 @@ async def main(): asyncio.run(main()) ``` - diff --git a/docs/index.md b/docs/index.md index 7a796d5..c146c82 100644 --- a/docs/index.md +++ b/docs/index.md @@ -7,7 +7,6 @@

-

@@ -20,7 +19,6 @@ Note: Please be aware that Agentic Security is designed as a safety scanner tool and not a foolproof solution. It cannot guarantee complete protection against all possible threats. - ## UI 🧙 booking-screen diff --git a/mkdocs.yml b/mkdocs.yml index 37bf108..da15002 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -34,9 +34,6 @@ plugins: python: paths: [agentic_security] -extra: - # hide the "Made with Material for MkDocs" message - generator: false footer: links: [] # Removes the default footer credits @@ -65,6 +62,7 @@ theme: favicon: "https://res.cloudinary.com/dq0w2rtm9/image/upload/v1737555066/r17hrkre246doczwmvbv.png" extra: + generator: false social: - icon: fontawesome/brands/github link: https://github.com/msoedov/agentic_security