Skip to content

Commit

Permalink
Add missing files
Browse files Browse the repository at this point in the history
  • Loading branch information
hupe1980 committed Apr 5, 2024
1 parent 928d353 commit 4f4c568
Show file tree
Hide file tree
Showing 4 changed files with 41 additions and 1 deletion.
3 changes: 2 additions & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -160,4 +160,5 @@ cython_debug/
#.idea/

.env
TODO.md
TODO.md
!aisploit/target
8 changes: 8 additions & 0 deletions aisploit/target/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
from .stdout import StdOutTarget
from .target import WrapperTarget, target

__all__ = [
"StdOutTarget",
"WrapperTarget",
"target",
]
15 changes: 15 additions & 0 deletions aisploit/target/stdout.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
import sys

from typing import IO


from ..core import BaseTarget


class StdOutTarget(BaseTarget):
def __init__(self, text_stream: IO[str] = sys.stdout) -> None:
self._text_stream = text_stream

def send_prompt(self, prompt: str) -> str:
self._text_stream.write(f"{prompt}\n")
return prompt
16 changes: 16 additions & 0 deletions aisploit/target/target.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
from typing import Callable


from ..core import BaseTarget


class WrapperTarget(BaseTarget):
def __init__(self, func: Callable) -> None:
self._func = func

def send_prompt(self, prompt: str) -> str:
return self._func(prompt)


def target(func: Callable) -> WrapperTarget:
return WrapperTarget(func)

0 comments on commit 4f4c568

Please sign in to comment.