diff --git a/.env b/.env index e69de29..11b953d 100644 --- a/.env +++ b/.env @@ -0,0 +1 @@ +INVARIANT_API_URL=https://explorer.invariantlabs.ai \ No newline at end of file diff --git a/proxy/routes/proxy.py b/proxy/routes/proxy.py index 2da82ad..b401dba 100644 --- a/proxy/routes/proxy.py +++ b/proxy/routes/proxy.py @@ -22,17 +22,20 @@ IGNORED_HEADERS = [ ] proxy = APIRouter() +MISSING_INVARIANT_AUTH_HEADER = "Missing invariant-authorization header" +MISSING_AUTH_HEADER = "Missing authorization header" +NOT_SUPPORTED_ENDPOINT = "Not supported OpenAI endpoint" +FAILED_TO_PUSH_TRACE = "Failed to push trace to the dataset" + def validate_headers( invariant_authorization: str = Header(None), authorization: str = Header(None) ): """Require the invariant-authorization and authorization headers to be present""" if invariant_authorization is None: - raise HTTPException( - status_code=400, detail="Missing invariant-authorization header" - ) + raise HTTPException(status_code=400, detail=MISSING_INVARIANT_AUTH_HEADER) if authorization is None: - raise HTTPException(status_code=400, detail="Missing authorization header") + raise HTTPException(status_code=400, detail=MISSING_AUTH_HEADER) @proxy.post( @@ -46,13 +49,11 @@ async def openai_proxy( ): """Proxy call to a language model provider""" if endpoint not in ALLOWED_OPEN_AI_ENDPOINTS: - raise HTTPException(status_code=404, detail="Not supported OpenAI endpoint") + raise HTTPException(status_code=404, detail=NOT_SUPPORTED_ENDPOINT) - headers = dict(request.headers) - - # Remove extra headers - for h in IGNORED_HEADERS: - headers.pop(h, None) + headers = { + k: v for k, v in request.headers.items() if k.lower() not in IGNORED_HEADERS + } headers["accept-encoding"] = "identity" request_body = await request.body() @@ -67,28 +68,22 @@ async def openai_proxy( response = await client.send(open_ai_request) try: json_response = response.json() + # push messages to the Invariant Explorer + # use both the request and response messages + messages = json.loads(request_body).get("messages", []) + messages += [ + choice["message"] for choice in json_response.get("choices", []) + ] _ = push_trace( dataset_name=dataset_name, - # Extract messages from the request body and response - messages=[ - json.loads(request_body)["messages"] - + [ - dict(choice["message"]) - for choice in json_response.get("choices", []) - ] - ], + messages=messages, invariant_authorization=request.headers.get("invariant-authorization"), ) except Exception as e: - raise HTTPException( - status_code=500, detail="Failed to push trace to the dataset" - ) from e + raise HTTPException(status_code=500, detail=FAILED_TO_PUSH_TRACE) from e # Detect if original request expects gzip encoding - original_accept_encoding = request.headers.get("accept-encoding", "") - should_gzip = "gzip" in original_accept_encoding.lower() - - if should_gzip: + if "gzip" in request.headers.get("accept-encoding", "").lower(): # Compress the response using gzip gzip_buffer = BytesIO() with gzip.GzipFile(mode="wb", fileobj=gzip_buffer) as gz: diff --git a/proxy/utils/__init__.py b/proxy/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/proxy/utils/explorer.py b/proxy/utils/explorer.py new file mode 100644 index 0000000..29c9f3b --- /dev/null +++ b/proxy/utils/explorer.py @@ -0,0 +1,40 @@ +"""Utility functions for the Invariant explorer.""" + +import os +from typing import Any, Dict, List + +from fastapi import HTTPException +from invariant_sdk.client import Client + +DEFAULT_API_URL = "https://explorer.invariantlabs.ai" + + +def push_trace( + messages: List[Dict[str, Any]], + dataset_name: str, + invariant_authorization: str, + api_url: str = DEFAULT_API_URL, +) -> Dict[str, str]: + """Pushes traces to the dataset on the Invariant Explorer. + + Args: + messages (List[Dict[str, Any]]): List of messages to push. + dataset_name (str): Name of the dataset. + invariant_authorization (str): Authorization token. + api_url (str): URL of the Invariant Explorer API. + + Returns: + Dict[str, str]: Response containing the trace ID. + """ + api_url = os.getenv("INVARIANT_API_URL", DEFAULT_API_URL) + api_key = invariant_authorization.split("Bearer ")[1] + client = Client(api_url=api_url, api_key=api_key) + try: + push_trace_response = client.create_request_and_push_trace( + messages=messages, dataset=dataset_name + ) + return {"trace_id": push_trace_response.id[0]} + except Exception as e: + raise HTTPException( + status_code=500, detail="Failed to push traces to the dataset" + ) from e