Merge branch 'main' into mcp-metadata

.github/workflows/tests_ci.yml

@@ -14,27 +14,47 @@ jobs:
     name: Build & Test
     runs-on: ubuntu-latest
     steps:
-        - name: Checkout
-          uses: actions/checkout@v4
+      - name: Checkout
+        uses: actions/checkout@v4
 
-        - name: Set Up Python
-          uses: actions/setup-python@v4
-          with:
-            python-version: "3.10"
+      - name: Set Up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.10"
 
-        - name: Install Dependencies
-          run: |
-            python -m pip install --upgrade pip
-            pip install pytest
+      - name: Install Dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install pytest
 
-        - name: Run unit tests
-          run: ./run.sh unit-tests -s -vv
-          continue-on-error: true
+      - name: Run unit tests
+        id: unit-tests
+        run: |
+          pip install -r tests/unit_tests/requirements.txt
+          ./run.sh unit-tests -s -vv
+        continue-on-error: true
 
-        - name: Run integration tests
-          env:
-            OPENAI_API_KEY: ${{ secrets.INVARIANT_TESTING_OPENAI_KEY }}
-            ANTHROPIC_API_KEY: ${{ secrets.INVARIANT_TESTING_ANTHROPIC_KEY}}
-            GEMINI_API_KEY: ${{ secrets.INVARIANT_TESTING_GEMINI_KEY }}
-            INVARIANT_API_KEY: ${{ secrets.INVARIANT_TESTING_GUARDRAILS_KEY }}
-          run: ./run.sh integration-tests -s -vv
+      - name: Run integration tests
+        id: integration-tests
+        env:
+          OPENAI_API_KEY: ${{ secrets.INVARIANT_TESTING_OPENAI_KEY }}
+          ANTHROPIC_API_KEY: ${{ secrets.INVARIANT_TESTING_ANTHROPIC_KEY }}
+          GEMINI_API_KEY: ${{ secrets.INVARIANT_TESTING_GEMINI_KEY }}
+          INVARIANT_API_KEY: ${{ secrets.INVARIANT_TESTING_GUARDRAILS_KEY }}
+        run: ./run.sh integration-tests -s -vv
+        continue-on-error: true
+
+      - name: Check test results
+        run: |
+          echo "Unit tests outcome: ${{ steps.unit-tests.outcome }}"
+          echo "Integration tests outcome: ${{ steps.integration-tests.outcome }}"
+          
+          if [[ "${{ steps.unit-tests.outcome }}" != "success" ]]; then
+            echo "Unit tests failed"
+            exit 1
+          fi
+
+          if [[ "${{ steps.integration-tests.outcome }}" != "success" ]]; then
+            echo "Integration tests failed"
+            exit 1
+          fi

gateway/routes/anthropic.py

@@ -14,10 +14,7 @@ from gateway.common.config_manager import (
     GatewayConfigManager,
     extract_guardrails_from_header,
 )
-from gateway.common.constants import (
-    CLIENT_TIMEOUT,
-    IGNORED_HEADERS,
-)
+from gateway.common.constants import CLIENT_TIMEOUT, IGNORED_HEADERS
 from gateway.common.guardrails import GuardrailAction, GuardrailRuleSet
 from gateway.common.request_context import RequestContext
 from gateway.converters.anthropic_to_invariant import (
@@ -383,6 +380,8 @@ class InstrumentedAnthropicStreamingResponse(InstrumentedStreamingResponse):
         # guardrailing response (if any)
         self.guardrails_execution_result = {}
 
+        self.sse_buffer = ""  # Buffer for incomplete events
+
     async def on_start(self):
         """Check guardrails in a pipelined fashion, before processing the first chunk (for input guardrailing)."""
         if self.context.guardrails:
@@ -437,16 +436,78 @@ class InstrumentedAnthropicStreamingResponse(InstrumentedStreamingResponse):
             yield chunk
 
     async def on_chunk(self, chunk):
-        """Process the chunk and update the merged_response"""
-        decoded_chunk = chunk.decode().strip()
-        if not decoded_chunk:
-            return
+        """
+        Process the chunk and update the merged_response.
+        Each chunk may contain multiple events, separated by double newlines.
+        Each event has type and data fields, separated by a newline.
+        It is possible that a chunk contains some incomplete events.
 
-        # process chunk and extend the merged_response
-        process_chunk(decoded_chunk, self.merged_response)
+        Example:
+
+        b'event: message_start\ndata: {"type":"message_start","message":
+        {"id":"msg_01LkayzAaw7b7QkUAw91psyx","type":"message","role":"assistant"
+        ,"model":"claude-3-5-sonnet-20241022","content":[],"stop_reason":null,
+        "stop_sequence":null,"usage":{"input_tokens":20,"cache_creation_input_to'
+
+        and
+
+        b'kens":0,"cache_read_input_tokens":0,"output_tokens":1}}}\n\nevent: content_block_start
+        \ndata: {"type":"content_block_start","index":0,"content_block"
+        :{"type":"text","text":""} }\n\nevent: ping
+        \ndata: {"type": "ping"}\n\nevent: content_block_delta
+        \ndata: {"type":"content_block_delta","index":0,"delta":{"type":
+        "text_delta","text":"Originally"} }\n\n'
+
+        In this case the first chunk ends with 'cache_creation_input_to' which is
+        continued in the next chunk.
+
+        in this case we need to maintain a buffer of the incomplete events.
+        We filter out the ping events and update a merged_response.
+        """
+        # Decode the chunk and add to buffer
+        decoded_chunk = chunk.decode("utf-8", errors="replace")
+        self.sse_buffer += decoded_chunk
+
+        # Process complete events from buffer
+        complete_events, incomplete_events = self.process_complete_events(
+            self.sse_buffer
+        )
+        self.sse_buffer = incomplete_events
+
+        # Check if we've received message_stop in any events
+        message_stop_received = False
+
+        # Update the merged_response based on complete events
+        for event in complete_events:
+            try:
+                if "event: message_stop" in event:
+                    message_stop_received = True
+
+                # Extract event data
+                lines = event.split("\n")
+                event_type = None
+                event_data = None
+
+                for line in lines:
+                    if line.startswith("event:"):
+                        event_type = line[6:].strip()
+                    elif line.startswith("data:"):
+                        event_data = line[5:].strip()
+
+                if event_data and event_type != "ping":  # Skip ping events
+                    try:
+                        event_json = json.loads(event_data)
+                        update_merged_response(event_json, self.merged_response)
+                    except json.JSONDecodeError as e:
+                        print(
+                            f"JSON parsing error in event: {e}. Event data: {event_data[:100]}...",
+                            flush=True,
+                        )
+            except Exception as e:
+                print(f"Error processing event: {e}", flush=True)
 
         # on last stream chunk, run output guardrails
-        if "event: message_stop" in decoded_chunk and self.context.guardrails:
+        if message_stop_received and self.context.guardrails:
             # Block on the guardrails check
             self.guardrails_execution_result = await get_guardrails_check_result(
                 self.context,
@@ -471,8 +532,32 @@ class InstrumentedAnthropicStreamingResponse(InstrumentedStreamingResponse):
                     value=f"event: error\ndata: {error_chunk}\n\n".encode()
                 )
 
+    def process_complete_events(self, buffer):
+        """Process the buffer and extract complete SSE events.
+
+        Returns:
+            Tuple[List[str], str]: A tuple containing a list of
+            complete events and the remaining buffer with incomplete events.
+        """
+        # Split on double newlines which separate SSE events
+        if not buffer:
+            return [], ""
+        events = []
+        remaining = buffer
+
+        # Process events that are complete (ending with \n\n)
+        while "\n\n" in remaining:
+            pos = remaining.find("\n\n")
+            if pos >= 0:
+                event = remaining[: pos + 2]
+                remaining = remaining[pos + 2 :]
+                if event.strip():  # Skip empty events
+                    events.append(event)
+
+        return events, remaining
+
     async def on_end(self):
-        """on_end: send full merged response to the exploree (if configured)"""
+        """on_end: send full merged response to the explorer (if configured)"""
         # don't block on the response from explorer (.create_task)
         if self.context.dataset_name:
             asyncio.create_task(
@@ -501,20 +586,6 @@ async def handle_streaming_response(
     )
 
 
-def process_chunk(chunk: str, merged_response: dict[str, Any]) -> None:
-    """
-    Process the chunk of text and update the merged_response
-    Example of chunk list can be find in:
-    ../../resources/streaming_chunk_text/anthropic.txt
-    """
-    for text_block in chunk.split("\n\n"):
-        # might be empty block
-        if len(text_block.split("\ndata:")) > 1:
-            event_text = text_block.split("\ndata:")[1]
-            event = json.loads(event_text)
-            update_merged_response(event, merged_response)
-
-
 def update_merged_response(
     event: dict[str, Any], merged_response: dict[str, Any]
 ) -> None:
@@ -530,22 +601,25 @@ def update_merged_response(
     final Message content array.
     3. One or more message_delta events, indicating top-level changes to the final Message object.
     A final message_stop event.
+    We filter out the ping eventss
 
     """
-    if event.get("type") == MESSAGE_START:
+    event_type = event.get("type")
+
+    if event_type == MESSAGE_START:
         merged_response.update(**event.get("message"))
-    elif event.get("type") == CONTENT_BLOCK_START:
+    elif event_type == CONTENT_BLOCK_START:
         index = event.get("index")
         if index >= len(merged_response.get("content")):
             merged_response["content"].append(event.get("content_block"))
         if event.get("content_block").get("type") == "tool_use":
             merged_response.get("content")[-1]["input"] = ""
-    elif event.get("type") == CONTENT_BLOCK_DELTA:
+    elif event_type == CONTENT_BLOCK_DELTA:
         index = event.get("index")
         delta = event.get("delta")
         if delta.get("type") == "text_delta":
             merged_response.get("content")[index]["text"] += delta.get("text")
         elif delta.get("type") == "input_json_delta":
             merged_response.get("content")[index]["input"] += delta.get("partial_json")
-    elif event.get("type") == MESSAGE_DELTA:
+    elif event_type == MESSAGE_DELTA:
         merged_response["usage"].update(**event.get("usage"))

run.sh

@@ -70,6 +70,10 @@ down() {
 unit_tests() {
   echo "Running unit tests..."
   PYTHONPATH=. pytest tests/unit_tests $@
+
+  TEST_EXIT_CODE=$?
+  echo "Unit tests finished with exit code: $TEST_EXIT_CODE"
+  return $TEST_EXIT_CODE
 }
 
 integration_tests() {
@@ -160,9 +164,13 @@ integration_tests() {
     -e INVARIANT_API_KEY="$INVARIANT_API_KEY" \
     --env-file ./tests/integration/.env.test \
     invariant-gateway-tests $@
+  TEST_EXIT_CODE=$?
+  echo "Integration tests finished with exit code: $TEST_EXIT_CODE"
 
   unset GATEWAY_ROOT_PATH
   unset GUARDRAILS_FILE_PATH
+
+  return $TEST_EXIT_CODE
 }
 
 # -----------------------------
@@ -185,10 +193,12 @@ case "$1" in
   "unit-tests")
     shift
     unit_tests $@
+    exit $?
     ;;
   "integration-tests")
     shift
     integration_tests $@
+    exit $?
     ;;
   *)
     echo "Usage: $0 [up|build|down|logs|unit-tests|integration-tests]"

tests/integration/anthropic/test_anthropic_with_tool_call.py

@@ -12,11 +12,10 @@ from typing import Dict, List
 # Add integration folder (parent) to sys.path
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
-from utils import get_anthropic_client
-
 import anthropic
 import pytest
 import requests
+from utils import get_anthropic_client
 
 # Pytest plugins
 pytest_plugins = ("pytest_asyncio",)
@@ -180,7 +179,7 @@ async def test_response_with_tool_call(explorer_api_url, gateway_url, push_to_ex
 
     weather_agent = WeatherAgent(gateway_url, push_to_explorer)
 
-    query = "Tell me the weather for New York"
+    query = "Tell me the weather for New York in Celsius"
 
     city = "new york"
     # Process each query
@@ -244,11 +243,12 @@ async def test_streaming_response_with_tool_call(
     """Test the chat completion with streaming for the weather agent."""
     weather_agent = WeatherAgent(gateway_url, push_to_explorer)
 
-    query = "Tell me the weather for New York"
+    query = "Tell me the weather for New York in Celsius"
     city = "new york"
 
     messages = [{"role": "user", "content": query}]
     response = weather_agent.get_streaming_response(messages)
+
     assert response is not None
     assert response[0][0].type == "text"
     assert response[0][1].type == "tool_use"
@@ -303,11 +303,11 @@ async def test_response_with_tool_call_with_image(
     """Test the chat completion with image for the weather agent."""
     weather_agent = WeatherAgent(gateway_url, push_to_explorer)
 
-    image_path = Path(__file__).parent.parent / "resources" / "images" / "new-york.jpeg"
+    image_path = Path(__file__).parent.parent / "resources" / "images" / "new-york.jpg"
 
     with image_path.open("rb") as image_file:
         base64_image = base64.b64encode(image_file.read()).decode("utf-8")
-        query = "get the weather in the city of this image"
+        query = "get the weather in the city of this image in Celsius"
         city = "new york"
         messages = [
             {
@@ -367,4 +367,3 @@ async def test_response_with_tool_call_with_image(
                 ].lower()
             )
             assert trace_messages[3]["role"] == "tool"
-            assert trace_messages[4]["role"] == "assistant"

tests/integration/anthropic/test_anthropic_without_tool_call.py

@@ -8,10 +8,9 @@ import uuid
 # Add integration folder (parent) to sys.path
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
-from utils import get_anthropic_client
-
 import pytest
 import requests
+from utils import get_anthropic_client
 
 # Pytest plugins
 pytest_plugins = ("pytest_asyncio",)
@@ -90,9 +89,9 @@ async def test_streaming_response_without_tool_call(
 
     cities = ["zurich", "new york", "london"]
     queries = [
-        "Can you introduce Zurich, Switzerland within 200 words?",
-        "Tell me the history of New York within 100 words?",
-        "How's the weather in London next week?",
+        "Can you introduce Zurich, Switzerland in 2 short sentences?",
+        "Tell me the history of New York within 2 short sentences.",
+        "Explain the geography of London in 2 short sentences.",
     ]
     # Process each query
     responses = []
@@ -102,7 +101,7 @@ async def test_streaming_response_without_tool_call(
 
         with client.messages.stream(
             model="claude-3-5-sonnet-20241022",
-            max_tokens=1024,
+            max_tokens=200,
             messages=messages,
         ) as response:
             for reply in response.text_stream:

tests/integration/guardrails/test_guardrails_anthropic.py

@@ -2,17 +2,16 @@
 
 import os
 import sys
-import uuid
 import time
+import uuid
 
 # Add integration folder (parent) to sys.path
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
-from utils import get_anthropic_client, create_dataset, add_guardrail_to_dataset
-
 import pytest
 import requests
 from anthropic import APIStatusError, BadRequestError
+from utils import add_guardrail_to_dataset, create_dataset, get_anthropic_client
 
 # Pytest plugins
 pytest_plugins = ("pytest_asyncio",)
@@ -164,7 +163,7 @@ async def test_tool_call_guardrail_from_file(
 
     if not do_stream:
         with pytest.raises(BadRequestError) as exc_info:
-            chat_response = client.messages.create(**request, stream=False)
+            _ = client.messages.create(**request, stream=False)
 
         assert exc_info.value.status_code == 400
         assert "[Invariant] The response did not pass the guardrails" in str(
@@ -174,10 +173,10 @@ async def test_tool_call_guardrail_from_file(
 
     else:
         with pytest.raises(APIStatusError) as exc_info:
-            chat_response = client.messages.create(**request, stream=True)
-
-            for _ in chat_response:
+            response = client.messages.create(**request, stream=True)
+            for _ in response:
                 pass
+
         assert (
             "[Invariant] The response did not pass the guardrails"
             in exc_info.value.message
@@ -535,10 +534,12 @@ async def test_preguardrailing_with_guardrails_from_explorer(
 
     else:
         if do_stream:
-            _ = client.messages.create(
+            response = client.messages.create(
                 **request,
                 stream=True,
             )
+            for _ in response:
+                pass
         else:
             _ = client.messages.create(
                 **request,

tests/integration/guardrails/test_guardrails_open_ai.py

@@ -2,17 +2,16 @@
 
 import os
 import sys
-import uuid
 import time
+import uuid
 
 # Add integration folder (parent) to sys.path
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
-from utils import get_open_ai_client, create_dataset, add_guardrail_to_dataset
-
 import pytest
 import requests
-from openai import BadRequestError, APIError
+from openai import APIError, BadRequestError
+from utils import add_guardrail_to_dataset, create_dataset, get_open_ai_client
 
 # Pytest plugins
 pytest_plugins = ("pytest_asyncio",)
@@ -532,10 +531,12 @@ async def test_preguardrailing_with_guardrails_from_explorer(
             assert "pun detected in user message" in str(exc_info.value)
     else:
         if do_stream:
-            _ = client.chat.completions.create(
+            response = client.chat.completions.create(
                 **request,
                 stream=True,
             )
+            for _ in response:
+                pass
         else:
             _ = client.chat.completions.create(
                 **request,

tests/integration/litellm/test_chat_without_tool_call.py

@@ -1,10 +1,12 @@
-import pytest
-import uuid
-from litellm import completion
-import litellm
-import time
-import requests
+"""Test the chat completions gateway calls with tool calling through litellm."""
+
 import os
+import time
+import uuid
+
+import pytest
+import requests
+from litellm import completion
 
 MODEL_API_KEYS = {
     "openai/gpt-4o": "OPENAI_API_KEY",
@@ -12,13 +14,21 @@ MODEL_API_KEYS = {
     "anthropic/claude-3-5-haiku-20241022": "ANTHROPIC_API_KEY",
 }
 
-@pytest.mark.parametrize("litellm_model", MODEL_API_KEYS.keys(),)
+
+@pytest.mark.parametrize(
+    "litellm_model",
+    MODEL_API_KEYS.keys(),
+)
 @pytest.mark.parametrize(
     "do_stream, push_to_explorer",
     [(False, False)],
 )
 async def test_chat_completion(
-    explorer_api_url: str, litellm_model: str, gateway_url: str, do_stream: bool, push_to_explorer: bool
+    explorer_api_url: str,
+    litellm_model: str,
+    gateway_url: str,
+    do_stream: bool,
+    push_to_explorer: bool,
 ):
     """Test the chat completions gateway calls with tool calling through litellm."""
     # Check if the API key is set in the environment variables
@@ -28,21 +38,23 @@ async def test_chat_completion(
     if not api_key:
         pytest.skip(f"Skipping {litellm_model} because {api_key_env_var} is not set")
 
-
     dataset_name = f"test-dataset-litellm-{litellm_model}-{uuid.uuid4()}"
-    base_url = (f"{gateway_url}/api/v1/gateway/{dataset_name}"
-    if push_to_explorer
-    else f"{gateway_url}/api/v1/gateway")
+    base_url = (
+        f"{gateway_url}/api/v1/gateway/{dataset_name}"
+        if push_to_explorer
+        else f"{gateway_url}/api/v1/gateway"
+    )
 
-    base_url += "/" + litellm_model.split("/")[0] #add provider name
+    base_url += "/" + litellm_model.split("/")[0]  # add provider name
     if litellm_model.split("/")[0] == "gemini":
-        base_url += f"/v1beta/models/{litellm_model.split('/')[1]}" #gemini expects the model name in the url
+        base_url += f"/v1beta/models/{litellm_model.split('/')[1]}"  # gemini expects the model name in the url
 
-    print(f"base_url: {base_url}")
     chat_response = completion(
         model=litellm_model,
         messages=[{"role": "user", "content": "What is the capital of France?"}],
-        extra_headers= {"Invariant-Authorization": "Bearer <some-key>"},
+        extra_headers={
+            "Invariant-Authorization": f"Bearer {os.environ['INVARIANT_API_KEY']}"
+        },
         stream=do_stream,
         base_url=base_url,
     )
@@ -92,4 +104,4 @@ async def test_chat_completion(
                 "role": "assistant",
                 "content": expected_assistant_message,
             },
-        ]
+        ]

tests/unit_tests/common/test_authorization.py

@@ -1,11 +1,12 @@
 """Tests for the authorization header extractor."""
 
 import os
-import sys
-from fastapi import HTTPException
 import random
 import string
+import sys
+
 import pytest
+from fastapi import HTTPException
 
 # Add root folder (parent) to sys.path
 sys.path.append(
@@ -14,14 +15,14 @@ sys.path.append(
     )
 )
 
-from gateway.common.config_manager import GatewayConfig
-from gateway.common.request_context import RequestContext
 from gateway.common.authorization import (
+    API_KEYS_SEPARATOR,
+    INVARIANT_AUTHORIZATION_HEADER,
     INVARIANT_GUARDRAIL_SERVICE_AUTHORIZATION_HEADER,
     extract_authorization_from_headers,
-    INVARIANT_AUTHORIZATION_HEADER,
-    API_KEYS_SEPARATOR,
 )
+from gateway.common.config_manager import GatewayConfig
+from gateway.common.request_context import RequestContext
 
 
 @pytest.mark.parametrize("push_to_explorer", [True, False])

tests/unit_tests/requirements.txt

@@ -0,0 +1,2 @@
+fastapi
+httpx