Fireworks AI Testing: Serverless Inference, Compound AI, and Function Calling

pip install fireworks-ai openai pytest pytest-asyncio httpx

import pytest
from unittest.mock import MagicMock, patch
import fireworks.client as fw_client
from openai import OpenAI  # Also works with Fireworks

FIREWORKS_BASE_URL = "https://api.fireworks.ai/inference/v1"

# Fireworks AI model identifiers
FIREWORKS_MODELS = {
    "llama-3-1-8b": "accounts/fireworks/models/llama-v3p1-8b-instruct",
    "llama-3-1-70b": "accounts/fireworks/models/llama-v3p1-70b-instruct",
    "mixtral-8x22b": "accounts/fireworks/models/mixtral-8x22b-instruct",
    "firefunction-v2": "accounts/fireworks/models/firefunction-v2",
    "qwen2-72b": "accounts/fireworks/models/qwen2-72b-instruct",
}

@pytest.fixture
def openai_fireworks_client():
    """OpenAI SDK configured for Fireworks AI."""
    return OpenAI(
        api_key="fw_test_key_not_real",
        base_url=FIREWORKS_BASE_URL
    )

@pytest.fixture
def mock_completion():
    """Standard mock completion response."""
    response = MagicMock()
    response.id = "chatcmpl-fireworks-abc123"
    response.model = "accounts/fireworks/models/llama-v3p1-8b-instruct"
    response.choices = [MagicMock()]
    response.choices[0].index = 0
    response.choices[0].message.role = "assistant"
    response.choices[0].message.content = "Test response"
    response.choices[0].message.tool_calls = None
    response.choices[0].finish_reason = "stop"
    response.usage = MagicMock()
    response.usage.prompt_tokens = 12
    response.usage.completion_tokens = 18
    response.usage.total_tokens = 30
    return response

@pytest.fixture
def function_call_response():
    """Mock response with function call."""
    tool_call = MagicMock()
    tool_call.id = "call_fw_abc123"
    tool_call.type = "function"
    tool_call.function.name = "search_products"
    tool_call.function.arguments = '{"query": "blue running shoes", "max_price": 150}'
    
    response = MagicMock()
    response.choices = [MagicMock()]
    response.choices[0].message.role = "assistant"
    response.choices[0].message.content = None
    response.choices[0].message.tool_calls = [tool_call]
    response.choices[0].finish_reason = "tool_calls"
    response.usage = MagicMock(prompt_tokens=80, completion_tokens=40, total_tokens=120)
    return response

def test_basic_completion(openai_fireworks_client, mock_completion):
    """Basic completion via Fireworks-compatible API."""
    with patch.object(openai_fireworks_client.chat.completions, "create",
                      return_value=mock_completion):
        response = openai_fireworks_client.chat.completions.create(
            model="accounts/fireworks/models/llama-v3p1-8b-instruct",
            messages=[{"role": "user", "content": "Hello"}]
        )
        
        assert response.choices[0].message.content == "Test response"
        assert response.usage.total_tokens == 30

def test_fireworks_model_id_format():
    """Fireworks model IDs should follow accounts/fireworks/models/ convention."""
    import re
    
    fireworks_model_pattern = re.compile(
        r'^accounts/[a-z0-9-]+/models/[a-z0-9._-]+$'
    )
    
    valid_models = [
        "accounts/fireworks/models/llama-v3p1-8b-instruct",
        "accounts/fireworks/models/firefunction-v2",
        "accounts/myorg/models/my-custom-model",
    ]
    
    invalid_models = [
        "llama-v3p1-8b-instruct",  # Missing prefix
        "accounts/fireworks/llama",  # Missing 'models/'
        "",
    ]
    
    for model in valid_models:
        assert fireworks_model_pattern.match(model), f"Should be valid: {model}"
    
    for model in invalid_models:
        assert not fireworks_model_pattern.match(model), f"Should be invalid: {model}"

def test_context_length_by_model():
    """Different Fireworks models have different context limits."""
    model_context_limits = {
        "accounts/fireworks/models/llama-v3p1-8b-instruct": 131072,
        "accounts/fireworks/models/llama-v3p1-70b-instruct": 131072,
        "accounts/fireworks/models/mixtral-8x22b-instruct": 65536,
        "accounts/fireworks/models/firefunction-v2": 8192,
    }
    
    def get_context_limit(model_id: str) -> int:
        return model_context_limits.get(model_id, 8192)  # Default fallback
    
    assert get_context_limit("accounts/fireworks/models/llama-v3p1-8b-instruct") == 131072
    assert get_context_limit("accounts/fireworks/models/firefunction-v2") == 8192
    assert get_context_limit("unknown/model") == 8192  # Fallback

def test_prompt_truncation_for_context_limit():
    """Messages exceeding context limit should be handled gracefully."""
    def estimate_tokens(text: str) -> int:
        """Rough token estimate (4 chars per token)."""
        return len(text) // 4
    
    def truncate_to_context(messages: list, model: str, max_tokens: int = 8192) -> list:
        """Truncate message history to fit context window."""
        total_tokens = sum(estimate_tokens(m.get("content", "")) for m in messages)
        
        while total_tokens > max_tokens and len(messages) > 1:
            # Remove oldest non-system message
            for i, msg in enumerate(messages):
                if msg["role"] != "system":
                    total_tokens -= estimate_tokens(msg.get("content", ""))
                    messages.pop(i)
                    break
        
        return messages
    
    # Build a long conversation
    messages = [{"role": "system", "content": "You are helpful."}]
    for i in range(100):
        messages.append({"role": "user", "content": f"Message {i}: " + "word " * 100})
        messages.append({"role": "assistant", "content": f"Response {i}"})
    
    truncated = truncate_to_context(messages, "firefunction-v2", max_tokens=2000)
    
    # Should have fewer messages but still have system prompt
    assert len(truncated) < 200
    assert truncated[0]["role"] == "system"

import json

def test_firefunction_tool_definition(openai_fireworks_client):
    """FireFunction tool definitions follow OpenAI format."""
    tools = [
        {
            "type": "function",
            "function": {
                "name": "search_products",
                "description": "Search for products in the catalog",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "Search query"
                        },
                        "max_price": {
                            "type": "number",
                            "description": "Maximum price in USD"
                        },
                        "category": {
                            "type": "string",
                            "enum": ["shoes", "clothing", "accessories"]
                        }
                    },
                    "required": ["query"]
                }
            }
        }
    ]
    
    with patch.object(openai_fireworks_client.chat.completions, "create") as mock:
        mock.return_value = MagicMock(
            choices=[MagicMock(
                message=MagicMock(content=None, tool_calls=[MagicMock()]),
                finish_reason="tool_calls"
            )],
            usage=MagicMock()
        )
        
        openai_fireworks_client.chat.completions.create(
            model="accounts/fireworks/models/firefunction-v2",
            messages=[{"role": "user", "content": "Find blue running shoes under $150"}],
            tools=tools
        )
        
        call_kwargs = mock.call_args.kwargs
        assert call_kwargs["tools"][0]["type"] == "function"
        assert call_kwargs["tools"][0]["function"]["name"] == "search_products"

def test_firefunction_response_parsing(openai_fireworks_client, function_call_response):
    """Function call responses should be correctly parsed."""
    with patch.object(openai_fireworks_client.chat.completions, "create",
                      return_value=function_call_response):
        response = openai_fireworks_client.chat.completions.create(
            model="accounts/fireworks/models/firefunction-v2",
            messages=[{"role": "user", "content": "Find blue running shoes"}],
            tools=[{"type": "function", "function": {"name": "search_products"}}]
        )
        
        assert response.choices[0].finish_reason == "tool_calls"
        
        tool_call = response.choices[0].message.tool_calls[0]
        assert tool_call.function.name == "search_products"
        
        args = json.loads(tool_call.function.arguments)
        assert "query" in args
        assert "blue running shoes" in args["query"].lower()

def test_parallel_function_calls(openai_fireworks_client):
    """FireFunction can call multiple functions in parallel."""
    tool_call_1 = MagicMock()
    tool_call_1.id = "call_001"
    tool_call_1.function.name = "get_weather"
    tool_call_1.function.arguments = '{"city": "Paris"}'
    
    tool_call_2 = MagicMock()
    tool_call_2.id = "call_002"
    tool_call_2.function.name = "get_exchange_rate"
    tool_call_2.function.arguments = '{"from": "USD", "to": "EUR"}'
    
    parallel_response = MagicMock()
    parallel_response.choices = [MagicMock()]
    parallel_response.choices[0].message.tool_calls = [tool_call_1, tool_call_2]
    parallel_response.choices[0].finish_reason = "tool_calls"
    
    with patch.object(openai_fireworks_client.chat.completions, "create",
                      return_value=parallel_response):
        response = openai_fireworks_client.chat.completions.create(
            model="accounts/fireworks/models/firefunction-v2",
            messages=[{
                "role": "user",
                "content": "What's the weather in Paris and current USD/EUR exchange rate?"
            }],
            tools=[
                {"type": "function", "function": {"name": "get_weather"}},
                {"type": "function", "function": {"name": "get_exchange_rate"}}
            ]
        )
        
        tool_calls = response.choices[0].message.tool_calls
        assert len(tool_calls) == 2
        
        function_names = {tc.function.name for tc in tool_calls}
        assert "get_weather" in function_names
        assert "get_exchange_rate" in function_names

def test_firefunction_structured_output(openai_fireworks_client):
    """FireFunction should produce schema-conforming JSON output."""
    from jsonschema import validate
    
    output_schema = {
        "type": "object",
        "properties": {
            "products": {
                "type": "array",
                "items": {
                    "type": "object",
                    "properties": {
                        "name": {"type": "string"},
                        "price": {"type": "number"},
                        "in_stock": {"type": "boolean"}
                    },
                    "required": ["name", "price", "in_stock"]
                }
            },
            "total_results": {"type": "integer"}
        },
        "required": ["products", "total_results"]
    }
    
    mock_json_output = {
        "products": [
            {"name": "Nike Air Zoom", "price": 129.99, "in_stock": True},
            {"name": "Adidas Ultraboost", "price": 149.99, "in_stock": False}
        ],
        "total_results": 2
    }
    
    mock_response = MagicMock()
    mock_response.choices = [MagicMock()]
    mock_response.choices[0].message.content = json.dumps(mock_json_output)
    mock_response.choices[0].finish_reason = "stop"
    
    with patch.object(openai_fireworks_client.chat.completions, "create",
                      return_value=mock_response):
        response = openai_fireworks_client.chat.completions.create(
            model="accounts/fireworks/models/firefunction-v2",
            messages=[{"role": "user", "content": "Search products"}],
            response_format={"type": "json_object"}
        )
        
        result = json.loads(response.choices[0].message.content)
        validate(instance=result, schema=output_schema)  # Raises if invalid
        
        assert len(result["products"]) == 2
        assert result["total_results"] == 2

import asyncio
from dataclasses import dataclass
from typing import Optional

@dataclass
class AgentStep:
    """Represents one step in a compound AI workflow."""
    model: str
    input: str
    output: Optional[str] = None
    tool_calls: Optional[list] = None
    error: Optional[str] = None

class CompoundAIWorkflow:
    """Multi-step AI workflow using Fireworks AI."""
    
    def __init__(self, client):
        self.client = client
        self.steps: list[AgentStep] = []
    
    def run_step(self, model: str, messages: list, tools=None) -> dict:
        kwargs = {"model": model, "messages": messages}
        if tools:
            kwargs["tools"] = tools
        
        response = self.client.chat.completions.create(**kwargs)
        
        step = AgentStep(model=model, input=messages[-1]["content"])
        
        if response.choices[0].message.tool_calls:
            step.tool_calls = response.choices[0].message.tool_calls
        else:
            step.output = response.choices[0].message.content
        
        self.steps.append(step)
        return response

def test_compound_workflow_step_tracking(openai_fireworks_client):
    """Compound AI workflows should track each step."""
    workflow = CompoundAIWorkflow(openai_fireworks_client)
    
    step_responses = []
    for i in range(3):
        mock_r = MagicMock()
        mock_r.choices = [MagicMock()]
        mock_r.choices[0].message.content = f"Step {i+1} output"
        mock_r.choices[0].message.tool_calls = None
        mock_r.choices[0].finish_reason = "stop"
        step_responses.append(mock_r)
    
    with patch.object(openai_fireworks_client.chat.completions, "create",
                      side_effect=step_responses):
        # Step 1: Extract intent
        workflow.run_step(
            model="accounts/fireworks/models/llama-v3p1-8b-instruct",
            messages=[{"role": "user", "content": "Book a flight to Paris"}]
        )
        
        # Step 2: Plan actions
        workflow.run_step(
            model="accounts/fireworks/models/firefunction-v2",
            messages=[{"role": "user", "content": "What tools do I need?"}]
        )
        
        # Step 3: Synthesize
        workflow.run_step(
            model="accounts/fireworks/models/llama-v3p1-70b-instruct",
            messages=[{"role": "user", "content": "Create the booking summary"}]
        )
    
    assert len(workflow.steps) == 3
    assert workflow.steps[0].model == "accounts/fireworks/models/llama-v3p1-8b-instruct"
    assert workflow.steps[2].model == "accounts/fireworks/models/llama-v3p1-70b-instruct"
    assert all(step.error is None for step in workflow.steps)

def test_compound_workflow_handles_partial_failure(openai_fireworks_client):
    """Workflow should handle step failures gracefully."""
    from openai import APIError
    
    success_response = MagicMock()
    success_response.choices = [MagicMock()]
    success_response.choices[0].message.content = "Success"
    success_response.choices[0].message.tool_calls = None
    
    failure = APIError("Model overloaded", response=MagicMock(status_code=503), body={})
    
    responses = [success_response, failure, success_response]
    call_count = [0]
    
    def side_effects(**kwargs):
        idx = call_count[0]
        call_count[0] += 1
        if isinstance(responses[idx], Exception):
            raise responses[idx]
        return responses[idx]
    
    results = []
    errors = []
    
    with patch.object(openai_fireworks_client.chat.completions, "create",
                      side_effect=side_effects):
        for i in range(3):
            try:
                r = openai_fireworks_client.chat.completions.create(
                    model="accounts/fireworks/models/llama-v3p1-8b-instruct",
                    messages=[{"role": "user", "content": f"Step {i}"}]
                )
                results.append(r)
            except APIError as e:
                errors.append(str(e))
    
    assert len(results) == 2
    assert len(errors) == 1
    assert "503" in errors[0] or "overloaded" in errors[0].lower()

def make_fireworks_stream(tokens: list[str]):
    """Create mock Fireworks streaming response."""
    for i, token in enumerate(tokens):
        chunk = MagicMock()
        chunk.id = "stream-fw-test"
        chunk.model = "accounts/fireworks/models/llama-v3p1-8b-instruct"
        chunk.choices = [MagicMock()]
        chunk.choices[0].delta.content = token
        chunk.choices[0].delta.role = "assistant" if i == 0 else None
        chunk.choices[0].finish_reason = "stop" if i == len(tokens) - 1 else None
        yield chunk

def test_streaming_response(openai_fireworks_client):
    """Streaming should yield tokens in order."""
    tokens = ["Fireworks", " AI", " is", " fast", "!"]
    
    with patch.object(openai_fireworks_client.chat.completions, "create",
                      return_value=make_fireworks_stream(tokens)):
        stream = openai_fireworks_client.chat.completions.create(
            model="accounts/fireworks/models/llama-v3p1-8b-instruct",
            messages=[{"role": "user", "content": "Tell me about Fireworks AI"}],
            stream=True
        )
        
        full_text = ""
        for chunk in stream:
            if chunk.choices[0].delta.content:
                full_text += chunk.choices[0].delta.content
        
        assert full_text == "Fireworks AI is fast!"

def test_streaming_latency_tracking(openai_fireworks_client):
    """Track time-to-first-token and total generation time."""
    import time
    
    def delayed_stream():
        time.sleep(0.05)  # Cold start simulation
        yield from make_fireworks_stream(["Token1", " Token2", " Token3"])
    
    with patch.object(openai_fireworks_client.chat.completions, "create",
                      return_value=delayed_stream()):
        start = time.monotonic()
        
        stream = openai_fireworks_client.chat.completions.create(
            model="accounts/fireworks/models/llama-v3p1-8b-instruct",
            messages=[{"role": "user", "content": "test"}],
            stream=True
        )
        
        first_token_time = None
        for chunk in stream:
            if chunk.choices[0].delta.content and first_token_time is None:
                first_token_time = time.monotonic() - start
        
        total_time = time.monotonic() - start
        
        assert first_token_time is not None
        assert total_time >= first_token_time
        # Production assertion: assert first_token_time < YOUR_SLA

from openai import AuthenticationError, RateLimitError, APIStatusError

def test_invalid_api_key_raises(openai_fireworks_client):
    """Bad API key should raise AuthenticationError."""
    with patch.object(openai_fireworks_client.chat.completions, "create",
                      side_effect=AuthenticationError(
                          message="Invalid API key",
                          response=MagicMock(status_code=401),
                          body={}
                      )):
        with pytest.raises(AuthenticationError):
            openai_fireworks_client.chat.completions.create(
                model="accounts/fireworks/models/llama-v3p1-8b-instruct",
                messages=[{"role": "user", "content": "test"}]
            )

def test_model_overloaded_503(openai_fireworks_client):
    """503 errors should trigger fallback to smaller model."""
    overloaded = APIStatusError(
        message="Model overloaded",
        response=MagicMock(status_code=503),
        body={"error": {"message": "Model overloaded, try again later"}}
    )
    
    fallback_response = MagicMock()
    fallback_response.choices = [MagicMock()]
    fallback_response.choices[0].message.content = "Fallback response"
    fallback_response.choices[0].message.tool_calls = None
    
    def fallback_completion(model, messages, **kwargs):
        if "70b" in model:
            raise overloaded
        return fallback_response
    
    with patch.object(openai_fireworks_client.chat.completions, "create",
                      side_effect=lambda **kwargs: fallback_completion(**kwargs)):
        def with_model_fallback(client, messages, preferred, fallback):
            try:
                return client.chat.completions.create(
                    model=preferred,
                    messages=messages
                )
            except APIStatusError as e:
                if e.response.status_code == 503:
                    return client.chat.completions.create(
                        model=fallback,
                        messages=messages
                    )
                raise
        
        result = with_model_fallback(
            openai_fireworks_client,
            messages=[{"role": "user", "content": "test"}],
            preferred="accounts/fireworks/models/llama-v3p1-70b-instruct",
            fallback="accounts/fireworks/models/llama-v3p1-8b-instruct"
        )
        
        assert result.choices[0].message.content == "Fallback response"

name: Fireworks AI Tests

on: [push, pull_request]

jobs:
  test:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      
      - uses: actions/setup-python@v5
        with:
          python-version: "3.11"
      
      - name: Install
        run: pip install fireworks-ai openai pytest pytest-asyncio pytest-cov jsonschema
      
      - name: Unit tests (mocked)
        run: pytest tests/ -v -m "not integration" --cov=src
        env:
          FIREWORKS_API_KEY: "fw_test_mock"
      
      - name: Model ID validation
        run: |
          python -c "
          import re
          pattern = re.compile(r'^accounts/[a-z0-9-]+/models/[a-z0-9._-]+$')
          
          with open('config/models.txt') as f:
              for line in f:
                  model_id = line.strip()
                  if model_id and not model_id.startswith('#'):
                      assert pattern.match(model_id), f'Invalid model ID: {model_id}'
          
          print('All model IDs valid')
          "
      
      - name: Integration tests (main only)
        if: github.ref == 'refs/heads/main'
        run: pytest tests/ -m integration -v
        env:
          FIREWORKS_API_KEY: ${{ secrets.FIREWORKS_API_KEY }}