diff --git a/proxy/routes/anthropic.py b/proxy/routes/anthropic.py index 86706f9..a9516ad 100644 --- a/proxy/routes/anthropic.py +++ b/proxy/routes/anthropic.py @@ -5,6 +5,7 @@ import json import httpx from typing import Any from utils.explorer import push_trace +from starlette.responses import StreamingResponse # from .open_ai import push_to_explorer proxy = APIRouter() @@ -20,6 +21,7 @@ IGNORED_HEADERS = [ "x-forwarded-proto", "x-forwarded-server", "x-real-ip", + "content-length" ] MISSING_INVARIANT_AUTH_HEADER = "Missing invariant-authorization header" @@ -32,6 +34,13 @@ END_REASONS = [ "stop_sequence" ] +MESSAGE_START = "message_start" +MESSGAE_DELTA = "message_delta" +MESSAGE_STOP = "message_stop" +CONTENT_BLOCK_START = "content_block_start" +CONTENT_BLOCK_DELTA = "content_block_delta" +CONTENT_BLOCK_STOP = "content_block_stop" + def validate_headers( invariant_authorization: str = Header(None), x_api_key: str = Header(None) ): @@ -56,6 +65,7 @@ async def anthropic_proxy( 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() @@ -63,18 +73,26 @@ async def anthropic_proxy( anthropic_url = f"https://api.anthropic.com/{endpoint}" client = httpx.AsyncClient() - + anthropic_request = client.build_request( "POST", anthropic_url, headers=headers, data=request_body ) - invariant_authorization = request.headers.get("invariant-authorization") - async with client: + response = await client.send(anthropic_request) + + if request_body_json.get("stream"): + return await handle_streaming_response(client, anthropic_request, dataset_name, invariant_authorization) + else: response = await client.send(anthropic_request) + if response.status_code != 200: + raise HTTPException( + status_code=response.status_code, + detail=f"Failed to fetch response: {response.text}", + ) await handle_non_streaming_response( response, dataset_name, request_body_json, invariant_authorization ) @@ -85,13 +103,14 @@ async def push_to_explorer( merged_response: dict[str, Any], request_body: dict[str, Any], invariant_authorization: str, + reformat: bool = True, ) -> None: """Pushes the full trace to the Invariant Explorer""" # Combine the messages from the request body and Anthropic response messages = request_body.get("messages", []) messages += [merged_response] - - messages = anthropic_to_invariant_messages(messages) + if reformat: + messages = anthropic_to_invariant_messages(messages) _ = await push_trace( dataset_name=dataset_name, messages=[messages], @@ -115,6 +134,83 @@ async def handle_non_streaming_response( invariant_authorization, ) +async def handle_streaming_response( + client: httpx.AsyncClient, + anthropic_request: httpx.Request, + dataset_name: str, + invariant_authorization: str +) -> StreamingResponse: + + format_invariant_response = [] + + async def event_generator() -> Any: + async with client.stream( + "POST", + anthropic_request.url, + headers=anthropic_request.headers, + content=anthropic_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(): + yield chunk + + process_chunk_text( + chunk, + format_invariant_response + ) + + if len(format_invariant_response) > 0 and format_invariant_response[-1].get("stop_reason") in END_REASONS: + await push_to_explorer( + dataset_name, + format_invariant_response[-1], + json.loads(anthropic_request.content), + invariant_authorization, + reformat = False + ) + + generator = event_generator() + + return StreamingResponse(generator, media_type="text/event-stream") + + +def process_chunk_text(chunk, format_invariant_response): + """ + Process the chunk of text and update the format_invariant_response + Example of chunk list: + chunk_list = [b'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_012KWB6kiKvzx7r1SKs5nGA1","type":"message","role":"assistant","model":"claude-3-5-sonnet-20241022","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":20,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":5}} }\n\nevent: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""} }\n\n' + ,b'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" making it an attractive destination for both business"} }\n\n' + , b'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" and leisure."} }\n\n' + , b'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}\n\n' + , b'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":301} }\n\n' + , b'event: message_stop\ndata: {"type":"message_stop" }\n\n'] + """ + + text_decode = chunk.decode().strip() + for text_block in text_decode.split("\n\n"): + text_data = text_block.split("\ndata:")[1] + text_json = json.loads(text_data) + update_format_invariant_response(text_json, format_invariant_response) + +def update_format_invariant_response(text_json, format_invariant_response): + if text_json.get("type") == MESSAGE_START: + message = text_json.get("message") + format_invariant_response.append({ + "id": message.get("id"), + "role": message.get("role"), + "content": "", + "model": message.get("model"), + "stop_reason": message.get("stop_reason"), + "stop_sequence": message.get("stop_sequence"), + }) + elif text_json.get("type") == CONTENT_BLOCK_DELTA and len(format_invariant_response) > 0: + format_invariant_response[-1]["content"] += text_json.get("delta").get("text") + elif text_json.get("type") == MESSGAE_DELTA and len(format_invariant_response) > 0: + format_invariant_response[-1]["stop_reason"] = text_json.get("delta").get("stop_reason") + def anthropic_to_invariant_messages( messages: list[dict], keep_empty_tool_response: bool = False ) -> list[dict]: diff --git a/tests/anthropic/test_claude_agent_streaming.py b/tests/anthropic/test_claude_agent_streaming.py new file mode 100644 index 0000000..0df1339 --- /dev/null +++ b/tests/anthropic/test_claude_agent_streaming.py @@ -0,0 +1,51 @@ +import anthropic +from typing import Dict +import os +from tavily import TavilyClient +import anthropic +from httpx import Client +import os +# from invariant import testing +tavily_client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY")) +import datetime + +def test_streaming_response_without_toolcall(): + # Example queries + dataset_name = "claude_streaming_agent_test" + str(datetime.datetime.now().strftime("%Y%m%d%H%M%S")) + invariant_api_key = os.environ.get("INVARIANT_API_KEY") + + client = anthropic.Anthropic( + http_client=Client( + headers={ + "Invariant-Authorization": f"Bearer {invariant_api_key}" + }, + ), + base_url=f"http://localhost/api/v1/proxy/{dataset_name}/anthropic", + ) + + cities = ["Zurich", "New York", "London"] + queries = [ + "Can you introduce Zurich city within 200 words?", + "Tell me the history of New York", + "How's the weather in London next week?" + ] + # Process each query + for index,query in enumerate(queries): + messages = [ + { + "role": "user", + "content": query + } + ] + response_text = "" + with client.messages.stream( + model="claude-3-5-sonnet-20241022", + max_tokens=1024, + messages=messages, + # stream = True + ) as response: + for reply in response.text_stream: + response_text += reply + print(reply, end="", flush=True) + assert reply != "" + assert cities[index] in response_text \ No newline at end of file