300 lines
9.9 KiB
Python
300 lines
9.9 KiB
Python
import json
|
|
import os
|
|
from typing import Any, AsyncIterator, Iterator, Optional
|
|
|
|
from httpx import AsyncClient, Client, HTTPStatusError
|
|
from guardrails_ai.types import (
|
|
Guard,
|
|
ValidationOutcome as IValidationOutcome,
|
|
)
|
|
|
|
from guardrails.errors import ValidationError
|
|
|
|
from guardrails.logger import logger
|
|
|
|
|
|
class GuardrailsApiClient:
|
|
ahttp_client: AsyncClient
|
|
http_client: Client
|
|
timeout: float
|
|
base_url: str
|
|
api_key: str
|
|
|
|
def __init__(self, base_url: Optional[str] = None, api_key: Optional[str] = None):
|
|
self.base_url = (
|
|
base_url
|
|
if base_url is not None
|
|
else os.environ.get("GUARDRAILS_BASE_URL", "http://localhost:8000")
|
|
)
|
|
self.api_key = (
|
|
api_key
|
|
if api_key is not None
|
|
else os.environ.get("GUARDRAILS_API_KEY", "x-guardrailsai-api-key")
|
|
)
|
|
self.timeout = 300
|
|
|
|
headers = {"x-guardrailsai-api-key": self.api_key}
|
|
|
|
self.ahttp_client = AsyncClient(
|
|
base_url=self.base_url, headers=headers, timeout=self.timeout
|
|
)
|
|
self.http_client = Client(
|
|
base_url=self.base_url, headers=headers, timeout=self.timeout
|
|
)
|
|
|
|
async def aupsert_guard(self, guard: Guard) -> Guard:
|
|
existing_guard: Guard | None = None
|
|
try:
|
|
existing_guard = await self.afetch_guard(guard.name)
|
|
except Exception:
|
|
pass
|
|
|
|
if existing_guard:
|
|
response = await self.ahttp_client.put(
|
|
f"/guards/{existing_guard.id}",
|
|
json=guard.model_dump(),
|
|
)
|
|
else:
|
|
response = await self.ahttp_client.post(
|
|
"/guards",
|
|
json=guard.model_dump(),
|
|
)
|
|
response.raise_for_status()
|
|
res_body = response.json()
|
|
return Guard.model_validate(res_body)
|
|
|
|
def upsert_guard(self, guard: Guard) -> Guard:
|
|
existing_guard: Guard | None = None
|
|
try:
|
|
existing_guard = self.fetch_guard(guard.name)
|
|
except Exception:
|
|
pass
|
|
|
|
if existing_guard:
|
|
response = self.http_client.put(
|
|
f"/guards/{existing_guard.id}",
|
|
json=guard.model_dump(),
|
|
)
|
|
else:
|
|
response = self.http_client.post(
|
|
"/guards",
|
|
json=guard.model_dump(),
|
|
)
|
|
response.raise_for_status()
|
|
res_body = response.json()
|
|
return Guard.model_validate(res_body)
|
|
|
|
async def afetch_guard(self, guard_name: str) -> Optional[Guard]:
|
|
try:
|
|
response = await self.ahttp_client.get(
|
|
f"/guards?name={guard_name}",
|
|
)
|
|
response.raise_for_status()
|
|
res_body = response.json()
|
|
first = res_body[0] if res_body and len(res_body) > 0 else None
|
|
if not first:
|
|
raise ValueError(f"No guard found for name {guard_name}")
|
|
return Guard.model_validate(first)
|
|
except Exception as e:
|
|
logger.error(f"Error fetching guard {guard_name}: {e}")
|
|
return None
|
|
|
|
def fetch_guard(self, guard_name: str) -> Optional[Guard]:
|
|
try:
|
|
response = self.http_client.get(
|
|
f"/guards?name={guard_name}",
|
|
)
|
|
response.raise_for_status()
|
|
res_body = response.json()
|
|
first = res_body[0] if res_body and len(res_body) > 0 else None
|
|
if not first:
|
|
raise ValueError(f"No guard found for name {guard_name}")
|
|
return Guard.model_validate(first)
|
|
except Exception as e:
|
|
logger.error(f"Error fetching guard {guard_name}: {e}")
|
|
return None
|
|
|
|
async def adelete_guard(self, guard_name: str) -> Optional[Guard]:
|
|
guard = await self.afetch_guard(guard_name)
|
|
if guard and guard.id:
|
|
response = await self.ahttp_client.delete(f"/guards/{guard.id}")
|
|
response.raise_for_status()
|
|
res_body = response.json()
|
|
return Guard.model_validate(res_body)
|
|
|
|
def delete_guard(self, guard_name: str) -> Optional[Guard]:
|
|
guard = self.fetch_guard(guard_name)
|
|
if guard and guard.id:
|
|
response = self.http_client.delete(f"/guards/{guard.id}")
|
|
response.raise_for_status()
|
|
res_body = response.json()
|
|
return Guard.model_validate(res_body)
|
|
|
|
async def avalidate(
|
|
self,
|
|
guard: Guard,
|
|
openai_api_key: Optional[str] = None,
|
|
*,
|
|
llm_output: Optional[str] = None,
|
|
num_reasks: Optional[int] = None,
|
|
prompt_params: Optional[dict[str, Any]] = None,
|
|
**kwargs,
|
|
):
|
|
try:
|
|
_openai_api_key = (
|
|
openai_api_key
|
|
if openai_api_key is not None
|
|
else os.environ.get("OPENAI_API_KEY")
|
|
)
|
|
headers = {}
|
|
if _openai_api_key:
|
|
headers = {"x-openai-api-key": _openai_api_key}
|
|
|
|
response = await self.ahttp_client.post(
|
|
f"/guards/{guard.id}/validate",
|
|
json={
|
|
"llm_output": llm_output,
|
|
"num_reasks": num_reasks,
|
|
"prompt_params": prompt_params,
|
|
**kwargs,
|
|
},
|
|
headers=headers,
|
|
)
|
|
response.raise_for_status()
|
|
res_body = response.json()
|
|
return IValidationOutcome.model_validate(res_body)
|
|
except HTTPStatusError as e:
|
|
if e.response.status_code == 400:
|
|
raise ValidationError(str(e)) from e
|
|
|
|
def validate(
|
|
self,
|
|
guard: Guard,
|
|
openai_api_key: Optional[str] = None,
|
|
*,
|
|
llm_output: Optional[str] = None,
|
|
num_reasks: Optional[int] = None,
|
|
prompt_params: Optional[dict[str, Any]] = None,
|
|
**kwargs,
|
|
):
|
|
try:
|
|
_openai_api_key = (
|
|
openai_api_key
|
|
if openai_api_key is not None
|
|
else os.environ.get("OPENAI_API_KEY")
|
|
)
|
|
headers = {}
|
|
if _openai_api_key:
|
|
headers = {"x-openai-api-key": _openai_api_key}
|
|
|
|
response = self.http_client.post(
|
|
f"/guards/{guard.id}/validate",
|
|
json={
|
|
"llm_output": llm_output,
|
|
"num_reasks": num_reasks,
|
|
"prompt_params": prompt_params,
|
|
**kwargs,
|
|
},
|
|
headers=headers,
|
|
)
|
|
response.raise_for_status()
|
|
res_body = response.json()
|
|
return IValidationOutcome.model_validate(res_body)
|
|
except HTTPStatusError as e:
|
|
if e.response.status_code == 400:
|
|
raise ValidationError(str(e)) from e
|
|
|
|
async def astream_validate(
|
|
self,
|
|
guard: Guard,
|
|
openai_api_key: Optional[str] = None,
|
|
*,
|
|
llm_output: Optional[str] = None,
|
|
num_reasks: Optional[int] = None,
|
|
prompt_params: Optional[dict[str, Any]] = None,
|
|
**kwargs,
|
|
) -> AsyncIterator[Any]:
|
|
_openai_api_key = (
|
|
openai_api_key
|
|
if openai_api_key is not None
|
|
else os.environ.get("OPENAI_API_KEY")
|
|
)
|
|
|
|
headers = {"Content-Type": "application/json"}
|
|
if _openai_api_key:
|
|
headers["x-openai-api-key"] = _openai_api_key
|
|
|
|
async with self.ahttp_client.stream(
|
|
"POST",
|
|
f"/guards/{guard.id}/validate",
|
|
json={
|
|
"llm_output": llm_output,
|
|
"num_reasks": num_reasks,
|
|
"prompt_params": prompt_params,
|
|
"stream": True,
|
|
**kwargs,
|
|
},
|
|
) as response:
|
|
if not response.is_success:
|
|
response.raise_for_status()
|
|
return
|
|
|
|
async for raw_chunk in response.aiter_text():
|
|
str_chunk = raw_chunk.strip()
|
|
if str_chunk:
|
|
str_chunk_data = "".join(str_chunk.split("\n"))
|
|
chunk = json.loads(str_chunk_data)
|
|
if chunk.get("error"):
|
|
raise Exception(chunk.get("error").get("message"))
|
|
yield IValidationOutcome.model_validate(chunk)
|
|
|
|
def stream_validate(
|
|
self,
|
|
guard: Guard,
|
|
openai_api_key: Optional[str] = None,
|
|
*,
|
|
llm_output: Optional[str] = None,
|
|
num_reasks: Optional[int] = None,
|
|
prompt_params: Optional[dict[str, Any]] = None,
|
|
**kwargs,
|
|
) -> Iterator[Any]:
|
|
_openai_api_key = (
|
|
openai_api_key
|
|
if openai_api_key is not None
|
|
else os.environ.get("OPENAI_API_KEY")
|
|
)
|
|
|
|
headers = {"Content-Type": "application/json"}
|
|
if _openai_api_key:
|
|
headers["x-openai-api-key"] = _openai_api_key
|
|
|
|
with self.http_client.stream(
|
|
"POST",
|
|
f"/guards/{guard.id}/validate",
|
|
json={
|
|
"llm_output": llm_output,
|
|
"num_reasks": num_reasks,
|
|
"prompt_params": prompt_params,
|
|
"stream": True,
|
|
**kwargs,
|
|
},
|
|
) as response:
|
|
if not response.is_success:
|
|
response.raise_for_status()
|
|
return
|
|
|
|
for raw_chunk in response.iter_text():
|
|
str_chunk = raw_chunk.strip()
|
|
if str_chunk:
|
|
str_chunk_data = "".join(str_chunk.split("\n"))
|
|
chunk = json.loads(str_chunk_data)
|
|
if chunk.get("error"):
|
|
raise Exception(chunk.get("error").get("message"))
|
|
yield IValidationOutcome.model_validate(chunk)
|
|
|
|
def get_history(self, guard_id: str, call_id: str) -> Any:
|
|
response = self.http_client.get(f"/guards/{guard_id}/history/{call_id}")
|
|
response.raise_for_status()
|
|
res_body = response.json()
|
|
return res_body
|