mirror of
https://github.com/invariantlabs-ai/invariant-gateway.git
synced 2026-07-09 12:27:50 +02:00
Merge branch 'main' into anthropic-implement
This commit is contained in:
Zishan
@@ -1,5 +1,6 @@
|
||||
services:
|
||||
explorer-proxy:
|
||||
container_name: explorer-proxy
|
||||
build:
|
||||
context: ./proxy
|
||||
dockerfile: ../proxy/Dockerfile.proxy
|
||||
@@ -8,6 +9,10 @@ services:
|
||||
- .env
|
||||
environment:
|
||||
- DEV_MODE=true
|
||||
volumes:
|
||||
- type: bind
|
||||
source: ./proxy
|
||||
target: /srv/proxy
|
||||
networks:
|
||||
- invariant-explorer-web
|
||||
ports: []
|
||||
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
+270
-34
@@ -1,9 +1,11 @@
|
||||
"""Proxy service to forward requests to the OpenAI APIs"""
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response
|
||||
from starlette.responses import StreamingResponse
|
||||
from utils.explorer import push_trace
|
||||
|
||||
ALLOWED_OPEN_AI_ENDPOINTS = {"chat/completions"}
|
||||
@@ -18,6 +20,7 @@ IGNORED_HEADERS = [
|
||||
"x-forwarded-server",
|
||||
"x-real-ip",
|
||||
]
|
||||
|
||||
proxy = APIRouter()
|
||||
|
||||
MISSING_INVARIANT_AUTH_HEADER = "Missing invariant-authorization header"
|
||||
@@ -44,7 +47,7 @@ async def openai_proxy(
|
||||
request: Request,
|
||||
dataset_name: str,
|
||||
endpoint: str,
|
||||
):
|
||||
) -> Response:
|
||||
"""Proxy calls to the OpenAI APIs"""
|
||||
if endpoint not in ALLOWED_OPEN_AI_ENDPOINTS:
|
||||
raise HTTPException(status_code=404, detail=NOT_SUPPORTED_ENDPOINT)
|
||||
@@ -54,43 +57,276 @@ async def openai_proxy(
|
||||
}
|
||||
headers["accept-encoding"] = "identity"
|
||||
|
||||
request_body = await request.body()
|
||||
request_body_bytes = await request.body()
|
||||
request_body_json = json.loads(request_body_bytes)
|
||||
|
||||
print("request_body", request_body)
|
||||
# Check if the request is for streaming
|
||||
is_streaming = request_body_json.get("stream", False)
|
||||
invariant_authorization = request.headers.get("invariant-authorization")
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
open_ai_request = client.build_request(
|
||||
"POST",
|
||||
f"https://api.openai.com/v1/{endpoint}",
|
||||
content=request_body,
|
||||
headers=headers,
|
||||
client = httpx.AsyncClient()
|
||||
open_ai_request = client.build_request(
|
||||
"POST",
|
||||
f"https://api.openai.com/v1/{endpoint}",
|
||||
content=request_body_bytes,
|
||||
headers=headers,
|
||||
)
|
||||
if is_streaming:
|
||||
return await stream_response(
|
||||
client,
|
||||
open_ai_request,
|
||||
dataset_name,
|
||||
request_body_json,
|
||||
invariant_authorization,
|
||||
)
|
||||
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,
|
||||
messages=[messages],
|
||||
invariant_authorization=request.headers.get("invariant-authorization"),
|
||||
else:
|
||||
async with client:
|
||||
response = await client.send(open_ai_request)
|
||||
return await handle_non_streaming_response(
|
||||
response, dataset_name, request_body_json, invariant_authorization
|
||||
)
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=500, detail=FAILED_TO_PUSH_TRACE + str(e)
|
||||
) from e
|
||||
|
||||
response_headers = dict(response.headers)
|
||||
response_headers.pop("Content-Encoding", None)
|
||||
response_headers.pop("Content-Length", None)
|
||||
|
||||
return Response(
|
||||
content=json.dumps(response.json()),
|
||||
status_code=response.status_code,
|
||||
media_type="application/json",
|
||||
headers=response_headers,
|
||||
async def stream_response(
|
||||
client: httpx.AsyncClient,
|
||||
open_ai_request: httpx.Request,
|
||||
dataset_name: str,
|
||||
request_body_json: dict[str, Any],
|
||||
invariant_authorization: str,
|
||||
) -> StreamingResponse:
|
||||
"""
|
||||
Handles streaming the OpenAI response to the client while building a merged_response
|
||||
The chunks are returned to the caller immediately
|
||||
The merged_response is built from the chunks as they are received
|
||||
It is sent to the Invariant Explorer at the end of the stream
|
||||
"""
|
||||
|
||||
async def event_generator() -> Any:
|
||||
# merged_response will be updated with the data from the chunks in the stream
|
||||
# At the end of the stream, this will be sent to the explorer
|
||||
merged_response = {
|
||||
"id": None,
|
||||
"object": "chat.completion",
|
||||
"created": None,
|
||||
"model": None,
|
||||
"choices": [],
|
||||
"usage": None,
|
||||
}
|
||||
# Each chunk in the stream contains a list called "choices" each entry in the list
|
||||
# has an index.
|
||||
# A choice has a field called "delta" which may contain a list called "tool_calls".
|
||||
# Maps the choice index in the stream to the index in the merged_response["choices"] list
|
||||
choice_mapping_by_index = {}
|
||||
# Combines the choice index and tool call index to uniquely identify a tool call
|
||||
tool_call_mapping_by_index = {}
|
||||
|
||||
async with client.stream(
|
||||
"POST",
|
||||
open_ai_request.url,
|
||||
headers=open_ai_request.headers,
|
||||
content=open_ai_request.content,
|
||||
) as response:
|
||||
if response.status_code != 200:
|
||||
yield json.dumps(
|
||||
{"error": f"Failed to fetch response: {response.status_code}"}
|
||||
).encode()
|
||||
return
|
||||
|
||||
async for chunk in response.aiter_bytes():
|
||||
chunk_text = chunk.decode().strip()
|
||||
if not chunk_text:
|
||||
continue
|
||||
|
||||
# Yield chunk immediately to the client (proxy behavior)
|
||||
yield chunk
|
||||
|
||||
# Process the chunk
|
||||
# This will update merged_response with the data from the chunk
|
||||
process_chunk_text(
|
||||
chunk_text,
|
||||
merged_response,
|
||||
choice_mapping_by_index,
|
||||
tool_call_mapping_by_index,
|
||||
)
|
||||
|
||||
# Send full merged response to the explorer
|
||||
await push_to_explorer(
|
||||
dataset_name,
|
||||
merged_response,
|
||||
request_body_json,
|
||||
invariant_authorization,
|
||||
)
|
||||
|
||||
return StreamingResponse(event_generator(), media_type="text/event-stream")
|
||||
|
||||
|
||||
def initialize_merged_response() -> dict[str, Any]:
|
||||
"""Initializes the full response dictionary"""
|
||||
return {
|
||||
"id": None,
|
||||
"object": "chat.completion",
|
||||
"created": None,
|
||||
"model": None,
|
||||
"choices": [],
|
||||
"usage": None,
|
||||
}
|
||||
|
||||
|
||||
def process_chunk_text(
|
||||
chunk_text: str,
|
||||
merged_response: dict[str, Any],
|
||||
choice_mapping_by_index: dict[int, int],
|
||||
tool_call_mapping_by_index: dict[str, dict[str, Any]],
|
||||
) -> None:
|
||||
"""Processes the chunk text and updates the merged_response to be sent to the explorer"""
|
||||
# Split the chunk text into individual JSON strings
|
||||
# A single chunk can contain multiple "data: " sections
|
||||
for json_string in chunk_text.split("\ndata: "):
|
||||
json_string = json_string.replace("data: ", "").strip()
|
||||
|
||||
if not json_string or json_string == "[DONE]":
|
||||
continue
|
||||
|
||||
try:
|
||||
json_chunk = json.loads(json_string)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
update_merged_response(
|
||||
json_chunk,
|
||||
merged_response,
|
||||
choice_mapping_by_index,
|
||||
tool_call_mapping_by_index,
|
||||
)
|
||||
|
||||
|
||||
def update_merged_response(
|
||||
json_chunk: dict[str, Any],
|
||||
merged_response: dict[str, Any],
|
||||
choice_mapping_by_index: dict[int, int],
|
||||
tool_call_mapping_by_index: dict[str, dict[str, Any]],
|
||||
) -> None:
|
||||
"""Updates the merged_response with the data (content, tool_calls, etc.) from the JSON chunk"""
|
||||
merged_response["id"] = merged_response["id"] or json_chunk.get("id")
|
||||
merged_response["created"] = merged_response["created"] or json_chunk.get("created")
|
||||
merged_response["model"] = merged_response["model"] or json_chunk.get("model")
|
||||
|
||||
for choice in json_chunk.get("choices", []):
|
||||
index = choice.get("index", 0)
|
||||
|
||||
if index not in choice_mapping_by_index:
|
||||
choice_mapping_by_index[index] = len(merged_response["choices"])
|
||||
merged_response["choices"].append(
|
||||
{
|
||||
"index": index,
|
||||
"message": {"role": "assistant"},
|
||||
"finish_reason": None,
|
||||
}
|
||||
)
|
||||
|
||||
existing_choice = merged_response["choices"][choice_mapping_by_index[index]]
|
||||
delta = choice.get("delta", {})
|
||||
|
||||
update_existing_choice_with_delta(
|
||||
existing_choice, delta, tool_call_mapping_by_index, choice_index=index
|
||||
)
|
||||
|
||||
|
||||
def update_existing_choice_with_delta(
|
||||
existing_choice: dict[str, Any],
|
||||
delta: dict[str, Any],
|
||||
tool_call_mapping_by_index: dict[str, dict[str, Any]],
|
||||
choice_index: int,
|
||||
) -> None:
|
||||
"""Updates the choice with the data from the delta"""
|
||||
content = delta.get("content")
|
||||
if content is not None:
|
||||
if "content" not in existing_choice["message"]:
|
||||
existing_choice["message"]["content"] = ""
|
||||
existing_choice["message"]["content"] += content
|
||||
|
||||
if isinstance(delta.get("tool_calls"), list):
|
||||
if "tool_calls" not in existing_choice["message"]:
|
||||
existing_choice["message"]["tool_calls"] = []
|
||||
|
||||
for tool in delta["tool_calls"]:
|
||||
tool_index = tool.get("index")
|
||||
tool_id = tool.get("id")
|
||||
name = tool.get("function", {}).get("name")
|
||||
arguments = tool.get("function", {}).get("arguments", "")
|
||||
|
||||
if tool_index is None:
|
||||
continue
|
||||
|
||||
choice_with_tool_call_index = f"{choice_index}-{tool_index}"
|
||||
|
||||
if choice_with_tool_call_index not in tool_call_mapping_by_index:
|
||||
tool_call_mapping_by_index[choice_with_tool_call_index] = {
|
||||
"index": tool_index,
|
||||
"id": tool_id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": name,
|
||||
"arguments": "",
|
||||
},
|
||||
}
|
||||
existing_choice["message"]["tool_calls"].append(
|
||||
tool_call_mapping_by_index[choice_with_tool_call_index]
|
||||
)
|
||||
|
||||
tool_call_entry = tool_call_mapping_by_index[choice_with_tool_call_index]
|
||||
|
||||
if tool_id:
|
||||
tool_call_entry["id"] = tool_id
|
||||
|
||||
if name:
|
||||
tool_call_entry["function"]["name"] = name
|
||||
|
||||
if arguments:
|
||||
tool_call_entry["function"]["arguments"] += arguments
|
||||
|
||||
finish_reason = delta.get("finish_reason")
|
||||
if finish_reason is not None:
|
||||
existing_choice["finish_reason"] = finish_reason
|
||||
|
||||
|
||||
async def push_to_explorer(
|
||||
dataset_name: str,
|
||||
merged_response: dict[str, Any],
|
||||
request_body: dict[str, Any],
|
||||
invariant_authorization: str,
|
||||
) -> None:
|
||||
"""Pushes the full trace to the Invariant Explorer"""
|
||||
# Combine the messages from the request body and the choices from the OpenAI response
|
||||
messages = request_body.get("messages", [])
|
||||
messages += [choice["message"] for choice in merged_response.get("choices", [])]
|
||||
|
||||
_ = await push_trace(
|
||||
dataset_name=dataset_name,
|
||||
messages=[messages],
|
||||
invariant_authorization=invariant_authorization,
|
||||
)
|
||||
|
||||
|
||||
async def handle_non_streaming_response(
|
||||
response: httpx.Response,
|
||||
dataset_name: str,
|
||||
request_body_json: dict[str, Any],
|
||||
invariant_authorization: str,
|
||||
):
|
||||
"""Handles non-streaming OpenAI responses"""
|
||||
json_response = response.json()
|
||||
await push_to_explorer(
|
||||
dataset_name, json_response, request_body_json, invariant_authorization
|
||||
)
|
||||
|
||||
response_headers = dict(response.headers)
|
||||
response_headers.pop("Content-Encoding", None)
|
||||
response_headers.pop("Content-Length", None)
|
||||
|
||||
return Response(
|
||||
content=json.dumps(json_response),
|
||||
status_code=response.status_code,
|
||||
media_type="application/json",
|
||||
headers=response_headers,
|
||||
)
|
||||
|
||||
Binary file not shown.
Binary file not shown.
+34
-16
@@ -3,35 +3,53 @@
|
||||
import os
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from fastapi import HTTPException
|
||||
from invariant_sdk.client import Client
|
||||
import httpx
|
||||
from invariant_sdk.types.push_traces import PushTracesRequest
|
||||
|
||||
DEFAULT_API_URL = "https://explorer.invariantlabs.ai"
|
||||
PUSH_ENDPOINT = "/api/v1/push/trace"
|
||||
|
||||
|
||||
def push_trace(
|
||||
messages: List[Dict[str, Any]],
|
||||
async def push_trace(
|
||||
messages: List[List[Dict[str, Any]]],
|
||||
dataset_name: str,
|
||||
invariant_authorization: str,
|
||||
) -> Dict[str, str]:
|
||||
"""Pushes traces to the dataset on the Invariant Explorer.
|
||||
|
||||
Args:
|
||||
messages (List[Dict[str, Any]]): List of messages to push.
|
||||
messages (List[List[Dict[str, Any]]]): List of messages to push.
|
||||
dataset_name (str): Name of the dataset.
|
||||
invariant_authorization (str): Authorization token.
|
||||
invariant_authorization (str): Authorization token from the
|
||||
invariant-authorization header.
|
||||
|
||||
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:
|
||||
# TODO: Change this to the async version once that is available
|
||||
push_trace_response = client.create_request_and_push_trace(
|
||||
messages=messages, dataset=dataset_name
|
||||
api_url = os.getenv("INVARIANT_API_URL", DEFAULT_API_URL).rstrip("/")
|
||||
# Remove any None values from the messages
|
||||
update_messages = [
|
||||
[{k: v for k, v in msg.items() if v is not None} for msg in msg_list]
|
||||
for msg_list in messages
|
||||
]
|
||||
request = PushTracesRequest(messages=update_messages, dataset=dataset_name)
|
||||
async with httpx.AsyncClient() as client:
|
||||
explorer_push_request = client.build_request(
|
||||
"POST",
|
||||
f"{api_url}{PUSH_ENDPOINT}",
|
||||
json=request.to_json(),
|
||||
headers={
|
||||
"Authorization": f"{invariant_authorization}",
|
||||
"Accept": "application/json",
|
||||
},
|
||||
)
|
||||
return {"trace_id": push_trace_response.id[0]}
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e)) from e
|
||||
try:
|
||||
response = await client.send(explorer_push_request)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except httpx.HTTPStatusError as e:
|
||||
print(f"Failed to push trace: {e.response.text}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
print(f"Unexpected error pushing trace: {str(e)}")
|
||||
return {"error": str(e)}
|
||||
|
||||
Reference in New Issue
Block a user