Skip to content

Commit 5171d46

Browse files
jmansdorferjmansdorfer
authored
PG-1058: adding responses and mcp components (#56)
* adding responses and mcp components * adding responses model var to actions * non-docs review changes * fixing test to account for empty mcp response when no mcp is available * upgrading deps
1 parent 46af446 commit 5171d46

File tree

16 files changed

+893
-37
lines changed

16 files changed

+893
-37
lines changed

‎.github/workflows/main.yml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ jobs:
2727
PREDICTIONGUARD_API_KEY: ${{ secrets.PREDICTIONGUARD_API_KEY }}
2828
PREDICTIONGUARD_URL: ${{ vars.PREDICTIONGUARD_URL }}
2929
TEST_CHAT_MODEL: ${{ vars.TEST_CHAT_MODEL }}
30+
TEST_RESPONSES_MODEL: ${{ vars.TEST_RESPONSES_MODEL }}
3031
TEST_TEXT_EMBEDDINGS_MODEL: ${{ vars.TEST_TEXT_EMBEDDINGS_MODEL }}
3132
TEST_MULTIMODAL_EMBEDDINGS_MODEL: ${{ vars.TEST_MULTIMODAL_EMBEDDINGS_MODEL }}
3233
TEST_VISION_MODEL: ${{ vars.TEST_VISION_MODEL }}

‎.github/workflows/pr.yml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ jobs:
3131
PREDICTIONGUARD_API_KEY: ${{ secrets.PREDICTIONGUARD_API_KEY }}
3232
PREDICTIONGUARD_URL: ${{ vars.PREDICTIONGUARD_URL }}
3333
TEST_CHAT_MODEL: ${{ vars.TEST_CHAT_MODEL }}
34+
TEST_RESPONSES_MODEL: ${{ vars.TEST_RESPONSES_MODEL }}
3435
TEST_TEXT_EMBEDDINGS_MODEL: ${{ vars.TEST_TEXT_EMBEDDINGS_MODEL }}
3536
TEST_MULTIMODAL_EMBEDDINGS_MODEL: ${{ vars.TEST_MULTIMODAL_EMBEDDINGS_MODEL }}
3637
TEST_VISION_MODEL: ${{ vars.TEST_VISION_MODEL }}

‎examples/responses.ipynb‎

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"metadata": {},
6+
"source": "## Using Responses with Prediction Guard"
7+
},
8+
{
9+
"cell_type": "markdown",
10+
"metadata": {},
11+
"source": [
12+
"### Set up"
13+
]
14+
},
15+
{
16+
"cell_type": "code",
17+
"execution_count": null,
18+
"metadata": {},
19+
"outputs": [],
20+
"source": [
21+
"# Import necessary packages\n",
22+
"import os\n",
23+
"import json\n",
24+
"\n",
25+
"from predictionguard import PredictionGuard\n",
26+
"\n",
27+
"\n",
28+
"# Set your Prediction Guard token and url as an environmental variable.\n",
29+
"os.environ[\"PREDICTIONGUARD_API_KEY\"] = \"<api key>\"\n",
30+
"os.environ[\"PREDICTIONGUARD_URL\"] = \"<url>\"\n",
31+
"\n",
32+
"# Or set your Prediction Guard token and url when initializing the PredictionGuard class.\n",
33+
"client = PredictionGuard(\n",
34+
" api_key=\"<api_key>\",\n",
35+
" url=\"<url>\"\n",
36+
")"
37+
]
38+
},
39+
{
40+
"cell_type": "markdown",
41+
"metadata": {},
42+
"source": "### Basic Responses"
43+
},
44+
{
45+
"cell_type": "code",
46+
"execution_count": null,
47+
"metadata": {},
48+
"outputs": [],
49+
"source": [
50+
"response = client.responses.create(\n",
51+
" model=\"Hermes-3-Llama-3.1-8B\",\n",
52+
" input=\"Tell me a funny joke about pirates.\"\n",
53+
")\n",
54+
"\n",
55+
"print(json.dumps(\n",
56+
" response,\n",
57+
" sort_keys=True,\n",
58+
" indent=4,\n",
59+
" separators=(',', ': ')\n",
60+
"))"
61+
]
62+
},
63+
{
64+
"cell_type": "markdown",
65+
"metadata": {},
66+
"source": "### Response with Images"
67+
},
68+
{
69+
"cell_type": "code",
70+
"execution_count": null,
71+
"metadata": {},
72+
"outputs": [],
73+
"source": [
74+
"input = [\n",
75+
" {\n",
76+
" \"role\": \"user\",\n",
77+
" \"content\": [\n",
78+
" {\n",
79+
" # Text to use for inference.\n",
80+
" \"type\": \"input_text\",\n",
81+
" \"text\": \"What's in this image?\"\n",
82+
" },\n",
83+
" {\n",
84+
" # Image to use for inference. Accepts image urls, files, and base64 encoded images, all under the \"image_url\" param.\n",
85+
" \"type\": \"input_image\",\n",
86+
" \"image_url\": \"https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg\"\n",
87+
" }\n",
88+
" ]\n",
89+
" },\n",
90+
"]\n",
91+
"\n",
92+
"image_response = client.responses.create(\n",
93+
" model=\"Qwen2.5-VL-7B-Instruct\",\n",
94+
" input=input\n",
95+
")\n",
96+
"\n",
97+
"print(json.dumps(\n",
98+
" image_response,\n",
99+
" sort_keys=True,\n",
100+
" indent=4,\n",
101+
" separators=(',', ': ')\n",
102+
"))"
103+
]
104+
},
105+
{
106+
"cell_type": "markdown",
107+
"metadata": {},
108+
"source": "### List Responses Models"
109+
},
110+
{
111+
"cell_type": "code",
112+
"execution_count": null,
113+
"metadata": {},
114+
"outputs": [],
115+
"source": [
116+
"model_list = client.responses.list_models()\n",
117+
"\n",
118+
"print(model_list)"
119+
]
120+
}
121+
],
122+
"metadata": {
123+
"language_info": {
124+
"name": "python"
125+
}
126+
},
127+
"nbformat": 4,
128+
"nbformat_minor": 2
129+
}

‎predictionguard/client.py‎

Lines changed: 17 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
from typing import Optional, Union
55

66
from .src.audio import Audio
7+
from .src.responses import Responses
78
from .src.chat import Chat
89
from .src.completions import Completions
910
from .src.detokenize import Detokenize
@@ -16,13 +17,15 @@
1617
from .src.toxicity import Toxicity
1718
from .src.pii import Pii
1819
from .src.injection import Injection
20+
from .src.mcp_servers import MCPServers
21+
from .src.mcp_tools import MCPTools
1922
from .src.models import Models
2023
from .version import __version__
2124

2225
__all__ = [
23-
"PredictionGuard", "Chat", "Completions", "Embeddings",
24-
"Audio", "Documents", "Rerank", "Tokenize", "Translate",
25-
"Detokenize", "Factuality", "Toxicity", "Pii", "Injection",
26+
"PredictionGuard", "Responses", "Chat", "Completions", "Embeddings",
27+
"Audio", "Documents", "Rerank", "Tokenize", "Translate", "Detokenize",
28+
"Factuality", "Toxicity", "Pii", "Injection", "MCPServers", "MCPTools",
2629
"Models"
2730
]
2831

@@ -81,11 +84,14 @@ def __init__(
8184
self._connect_client()
8285

8386
# Pass Prediction Guard class variables to inner classes
87+
self.responses: Responses = Responses(self.api_key, self.url, self.timeout)
88+
"""Responses allows for the usage of LLMs intended for agentic usages."""
89+
8490
self.chat: Chat = Chat(self.api_key, self.url, self.timeout)
85-
"""Chat generates chat completions based on a conversation history"""
91+
"""Chat generates chat completions based on a conversation history."""
8692

8793
self.completions: Completions = Completions(self.api_key, self.url, self.timeout)
88-
"""Completions generates text completions based on the provided input"""
94+
"""Completions generates text completions based on the provided input."""
8995

9096
self.embeddings: Embeddings = Embeddings(self.api_key, self.url, self.timeout)
9197
"""Embedding generates chat completions based on a conversation history."""
@@ -120,6 +126,12 @@ def __init__(
120126
self.detokenize: Detokenize = Detokenize(self.api_key, self.url, self.timeout)
121127
"""Detokenizes generates text for input tokens."""
122128

129+
self.mcp_servers: MCPServers = MCPServers(self.api_key, self.url, self.timeout)
130+
"""MCPServers lists all the MCP servers available in the Prediction Guard API."""
131+
132+
self.mcp_tools: MCPTools = MCPTools(self.api_key, self.url, self.timeout)
133+
"""MCPTools lists all the MCP tools available in the Prediction Guard API."""
134+
123135
self.models: Models = Models(self.api_key, self.url, self.timeout)
124136
"""Models lists all of the models available in the Prediction Guard API."""
125137

‎predictionguard/src/chat.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -122,7 +122,7 @@ def create(
122122
:param model: The ID(s) of the model to use.
123123
:param messages: The content of the call, an array of dictionaries containing a role and content.
124124
:param input: A dictionary containing the PII and injection arguments.
125-
:param output: A dictionary containing the consistency, factuality, and toxicity arguments.
125+
:param output: A dictionary containing the factuality, and toxicity arguments.
126126
:param frequency_penalty: The frequency penalty to use.
127127
:param logit_bias: The logit bias to use.
128128
:param max_completion_tokens: The maximum amount of tokens the model should return.

‎predictionguard/src/completions.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ def create(
7171
:param model: The ID(s) of the model to use.
7272
:param prompt: The prompt(s) to generate completions for.
7373
:param input: A dictionary containing the PII and injection arguments.
74-
:param output: A dictionary containing the consistency, factuality, and toxicity arguments.
74+
:param output: A dictionary containing the factuality, and toxicity arguments.
7575
:param echo: A boolean indicating whether to echo the prompt(s) to the output.
7676
:param frequency_penalty: The frequency penalty to use.
7777
:param logit_bias: The logit bias to use.

‎predictionguard/src/mcp_servers.py‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
import requests
2+
from typing import Any, Dict, Optional
3+
4+
from ..version import __version__
5+
6+
7+
class MCPServers:
8+
"""
9+
MCPServers lists all the MCP servers available in the Prediction Guard API.
10+
11+
Usage::
12+
13+
import os
14+
import json
15+
16+
from predictionguard import PredictionGuard
17+
18+
# Set your Prediction Guard token and url as an environmental variable.
19+
os.environ["PREDICTIONGUARD_API_KEY"] = "<api key>"
20+
os.environ["PREDICTIONGUARD_URL"] = "<url>"
21+
22+
# Or set your Prediction Guard token and url when initializing the PredictionGuard class.
23+
client = PredictionGuard(
24+
api_key="<api_key>",
25+
url="<url>"
26+
)
27+
28+
response = client.mcp_servers.list()
29+
30+
print(json.dumps(
31+
response,
32+
sort_keys=True,
33+
indent=4,
34+
separators=(",", ": ")
35+
))
36+
"""
37+
38+
def __init__(self, api_key, url, timeout):
39+
self.api_key = api_key
40+
self.url = url
41+
self.timeout = timeout
42+
43+
def list(self) -> Dict[str, Any]:
44+
"""
45+
Creates a mcp_servers list request in the Prediction Guard REST API.
46+
47+
:return: A dictionary containing the metadata of all the MCP servers.
48+
"""
49+
50+
# Run _list_mcp_servers
51+
choices = self._list_mcp_servers()
52+
return choices
53+
54+
def _list_mcp_servers(self):
55+
"""
56+
Function to list available MCP servers.
57+
"""
58+
59+
headers = {
60+
"Content-Type": "application/json",
61+
"Authorization": "Bearer " + self.api_key,
62+
"User-Agent": "Prediction Guard Python Client: " + __version__,
63+
}
64+
65+
response = requests.request(
66+
"GET", self.url + "/mcp_servers", headers=headers, timeout=self.timeout
67+
)
68+
69+
if response.status_code == 200:
70+
ret = response.json()
71+
return ret
72+
elif response.status_code == 429:
73+
raise ValueError(
74+
"Could not connect to Prediction Guard API. "
75+
"Too many requests, rate limit or quota exceeded."
76+
)
77+
else:
78+
# Check if there is a JSON body in the response. Read that in,
79+
# print out the error field in the JSON body, and raise an exception.
80+
err = ""
81+
try:
82+
err = response.json()["error"]
83+
except Exception:
84+
pass
85+
raise ValueError("Could not check for injection. " + err)

‎predictionguard/src/mcp_tools.py‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
import requests
2+
from typing import Any, Dict, Optional
3+
4+
from ..version import __version__
5+
6+
7+
class MCPTools:
8+
"""
9+
MCPTools lists all the MCP tools available in the Prediction Guard API.
10+
11+
Usage::
12+
13+
import os
14+
import json
15+
16+
from predictionguard import PredictionGuard
17+
18+
# Set your Prediction Guard token and url as an environmental variable.
19+
os.environ["PREDICTIONGUARD_API_KEY"] = "<api key>"
20+
os.environ["PREDICTIONGUARD_URL"] = "<url>"
21+
22+
# Or set your Prediction Guard token and url when initializing the PredictionGuard class.
23+
client = PredictionGuard(
24+
api_key="<api_key>",
25+
url="<url>"
26+
)
27+
28+
response = client.mcp_tools.list()
29+
30+
print(json.dumps(
31+
response,
32+
sort_keys=True,
33+
indent=4,
34+
separators=(",", ": ")
35+
))
36+
"""
37+
38+
def __init__(self, api_key, url, timeout):
39+
self.api_key = api_key
40+
self.url = url
41+
self.timeout = timeout
42+
43+
def list(self) -> Dict[str, Any]:
44+
"""
45+
Creates a mcp_tools list request in the Prediction Guard REST API.
46+
47+
:return: A dictionary containing the metadata of all the MCP tools.
48+
"""
49+
50+
# Run _list_mcp_tools
51+
choices = self._list_mcp_tools()
52+
return choices
53+
54+
def _list_mcp_tools(self):
55+
"""
56+
Function to list available MCP tools.
57+
"""
58+
59+
headers = {
60+
"Content-Type": "application/json",
61+
"Authorization": "Bearer " + self.api_key,
62+
"User-Agent": "Prediction Guard Python Client: " + __version__,
63+
}
64+
65+
response = requests.request(
66+
"GET", self.url + "/mcp_tools", headers=headers, timeout=self.timeout
67+
)
68+
69+
if response.status_code == 200:
70+
ret = response.json()
71+
return ret
72+
elif response.status_code == 429:
73+
raise ValueError(
74+
"Could not connect to Prediction Guard API. "
75+
"Too many requests, rate limit or quota exceeded."
76+
)
77+
else:
78+
# Check if there is a JSON body in the response. Read that in,
79+
# print out the error field in the JSON body, and raise an exception.
80+
err = ""
81+
try:
82+
err = response.json()["error"]
83+
except Exception:
84+
pass
85+
raise ValueError("Could not check for injection. " + err)

0 commit comments

Comments
 (0)