Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

FEAT: Added support to target an Ollama endpoint as a prompt chat target #141

Merged
merged 6 commits into from
Apr 14, 2024

Conversation

uskr
Copy link
Contributor

@uskr uskr commented Apr 11, 2024

Description

This is a basic implementation for a prompt target for Ollama endpoint. This allows to use all kinds of models locally. A few benefits here are it allows to get it started with pyRIT a lot easier, no cloud environment required. Also allows for testing models locally without making them available online/sharing.

Tests and Documentation

Example of how to use the prompt target:

import os
import uuid

from pyrit.prompt_target import OllamaChatTarget

with OllamaChatTarget(
    endpoint_uri="http://127.0.0.1:11434/api/chat",
    model_name="mistral",
    ) as target_llm:

    prompt = "test"
    print(target_llm.send_prompt(normalized_prompt=prompt, conversation_id=str(uuid.uuid4()), normalizer_id=None))

@uskr
Copy link
Contributor Author

uskr commented Apr 11, 2024 via email

@romanlutz
Copy link
Contributor

Thank you for the contribution @uskr ! This is awesome!

Can we add unit tests where the response is mocked? I can point you to examples if you don't know what I mean. Just to ensure the logic doesn't break over time as it or currently used in notebooks (which current function as our e2e tests).

@uskr
Copy link
Contributor Author

uskr commented Apr 13, 2024

Thank you for the contribution @uskr ! This is awesome!

Can we add unit tests where the response is mocked? I can point you to examples if you don't know what I mean. Just to ensure the logic doesn't break over time as it or currently used in notebooks (which current function as our e2e tests).

Done!

Copy link
Contributor

@romanlutz romanlutz left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for adding the tests! Looks great to me! @rlundeen2 @cseifert1 @rdheekonda wdyt?

@uskr
Copy link
Contributor Author

uskr commented Apr 13, 2024

Do you want me to remove the 3 unused imports so the build_and_test job pass?

@romanlutz
Copy link
Contributor

Do you want me to remove the 3 unused imports so the build_and_test job pass?

Yes! You can run the same tests locally with pre-commit run --all-files

@uskr
Copy link
Contributor Author

uskr commented Apr 14, 2024

Done.

@rlundeen2 rlundeen2 merged commit 334861b into Azure:main Apr 14, 2024
4 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

3 participants