feat(Add pydantic-ai agent):

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())

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.

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())
 ```
-

docs/index.md

@@ -7,7 +7,6 @@
     <br />
     <br />
 
-
 </p>
 </p>
 
@@ -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 🧙
 
 <img width="100%" alt="booking-screen" src="https://res.cloudinary.com/dq0w2rtm9/image/upload/v1736433557/z0bsyzhsqlgcr3w4ovwp.gif">

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