eval-protocol 0.0.3__py3-none-any.whl

eval_protocol/generation/clients.py

@@ -0,0 +1,248 @@
+"""
+Model clients for generating responses from various LLM APIs.
+"""
+
+import abc
+import asyncio
+import json  # For parsing content as JSON
+import logging
+import uuid  # For generating tool call IDs if not provided
+from typing import Any, Dict, List, Optional
+
+import aiohttp
+from omegaconf import DictConfig
+from pydantic import BaseModel, Field  # Added for new models
+
+logger = logging.getLogger(__name__)
+
+
+# Pydantic models for structured tool calls and generation results
+class ToolCallFunction(BaseModel):
+    name: str
+    arguments: str  # Should be a JSON string
+
+
+class ToolCall(BaseModel):
+    id: str
+    type: str = "function"  # OpenAI default
+    function: ToolCallFunction
+
+
+class GenerationResult(BaseModel):
+    content: Optional[str] = None
+    tool_calls: Optional[List[ToolCall]] = None
+
+
+class ModelClient(abc.ABC):
+    """Abstract base class for model clients."""
+
+    def __init__(self, client_config: DictConfig):
+        self.client_config = client_config
+        self.model_name = client_config.get("model_name", "unknown_model")
+        self.temperature = client_config.get("temperature", 0.0)
+        self.max_tokens = client_config.get("max_tokens", 1024)
+        self.top_p = client_config.get("top_p", 0.95)
+        self.top_k = client_config.get("top_k", 20)
+        self.min_p = client_config.get("min_p", 0.0)
+        # Add reasoning_effort, defaulting to None if not specified in config
+        self.reasoning_effort = client_config.get("reasoning_effort", None)
+
+    @abc.abstractmethod
+    async def generate(
+        self,
+        messages: List[Dict[str, str]],
+        session: aiohttp.ClientSession,
+        tools: Optional[List[Dict[str, Any]]] = None,  # Added tools parameter
+    ) -> GenerationResult:  # Changed return type
+        """Generates a response from the model given a list of messages."""
+        pass
+
+
+class FireworksModelClient(ModelClient):
+    """Client for Fireworks AI models."""
+
+    def __init__(self, client_config: DictConfig, api_key: str):
+        super().__init__(client_config)
+        self.api_key = api_key
+        self.api_base = client_config.get("api_base", "https://api.fireworks.ai/inference/v1")
+        # TODO: Initialize rate limiter, retry policy from client_config.api_params
+
+    async def generate(
+        self,
+        messages: List[Dict[str, str]],
+        session: aiohttp.ClientSession,
+        tools: Optional[List[Dict[str, Any]]] = None,
+    ) -> GenerationResult:
+        url = f"{self.api_base}/chat/completions"
+
+        payload: Dict[str, Any] = {
+            "model": self.model_name,
+            "messages": messages,
+            "temperature": self.temperature,
+            "max_tokens": self.max_tokens,
+        }
+        if self.top_p is not None:
+            payload["top_p"] = self.top_p
+
+        if tools:
+            payload["tools"] = tools
+            # Fireworks API might use "function" or "any" or specific tool name for tool_choice.
+            # "auto" is common for OpenAI. If Fireworks needs specific, this might need adjustment.
+            # Or, if it's like older OpenAI, it might not use tool_choice if tools are present.
+            # For now, let's assume "auto" or that it's implicit if "tools" is provided.
+            # The user's log shows the LLM is attempting tool calls even with the simpler prompt,
+            # implying the `tools` parameter is having an effect or the model is well-primed.
+            payload["tool_choice"] = "auto"
+
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+            "Accept": "application/json",
+        }
+
+        debug_payload_log = json.loads(json.dumps(payload))
+        if "messages" in debug_payload_log and debug_payload_log["messages"]:
+            if debug_payload_log["messages"][-1].get("content"):  # Check if content exists
+                debug_payload_log["messages"][-1]["content"] = (
+                    str(debug_payload_log["messages"][-1]["content"])[:50] + "..."
+                )
+        logger.debug(f"Calling Fireworks API: {url}, Payload: {debug_payload_log}")
+
+        try:
+            for attempt in range(self.client_config.get("api_params", {}).get("max_retries", 3) + 1):
+                async with session.post(url, json=payload, headers=headers) as response:
+                    if response.status == 200:
+                        data = await response.json()
+                        if data.get("choices") and len(data["choices"]) > 0:
+                            choice = data["choices"][0]
+                            message = choice.get("message", {})
+
+                            # 1. Check for native OpenAI-style tool_calls field
+                            if message.get("tool_calls"):
+                                tool_calls_data = message["tool_calls"]
+                                parsed_tool_calls = []
+                                for tc_data in tool_calls_data:
+                                    if tc_data.get("type") == "function" and tc_data.get("function"):
+                                        parsed_tool_calls.append(
+                                            ToolCall(
+                                                id=tc_data.get(
+                                                    "id", f"call_{uuid.uuid4().hex[:8]}"
+                                                ),  # Generate ID if missing
+                                                type="function",
+                                                function=ToolCallFunction(
+                                                    name=tc_data["function"]["name"],
+                                                    arguments=tc_data["function"]["arguments"],
+                                                ),
+                                            )
+                                        )
+                                if parsed_tool_calls:
+                                    logger.debug(f"Parsed native tool_calls: {parsed_tool_calls}")
+                                    return GenerationResult(tool_calls=parsed_tool_calls)
+
+                            # 2. If no native tool_calls, check if content is a JSON string representing a tool call
+                            # This handles the case where the LLM puts the tool call JSON into the content field.
+                            # The user's log shows content like: "{\"type\": \"function\", \"name\": \"move_file\", ...}"
+                            if message.get("content"):
+                                content_str = message["content"]
+                                try:
+                                    # Attempt to parse content as JSON
+                                    potential_tool_call_data = json.loads(content_str)
+
+                                    # Check if it matches the OpenAI tool call structure (single call in content)
+                                    # e.g., {"type": "function", "function": {"name": "...", "arguments": "{...}"}}
+                                    # or the structure the LLM actually produced: {"type": "function", "name": "...", "parameters": {...}}
+
+                                    parsed_tool_calls_from_content = []
+                                    # Handle if content is a list of tool calls (less likely but possible)
+                                    if isinstance(potential_tool_call_data, list):
+                                        data_to_check = potential_tool_call_data
+                                    else:  # Assume it's a single tool call object
+                                        data_to_check = [potential_tool_call_data]
+
+                                    for item in data_to_check:
+                                        if isinstance(item, dict) and item.get("type") == "function":
+                                            func_details = item.get("function")  # OpenAI style
+                                            if func_details and "name" in func_details and "arguments" in func_details:
+                                                parsed_tool_calls_from_content.append(
+                                                    ToolCall(
+                                                        id=item.get(
+                                                            "id",
+                                                            f"call_{uuid.uuid4().hex[:8]}",
+                                                        ),
+                                                        type="function",
+                                                        function=ToolCallFunction(
+                                                            name=func_details["name"],
+                                                            arguments=func_details["arguments"],
+                                                        ),
+                                                    )
+                                                )
+                                                continue  # Found valid OpenAI style tool call
+
+                                            # Check for the LLM's observed output format: {"type": "function", "name": ..., "parameters": ...}
+                                            # This is slightly different from OpenAI's `function.arguments` being a string.
+                                            # Here, `parameters` is an object. We need to dump it to string for `ToolCallFunction.arguments`.
+                                            llm_name = item.get("name")
+                                            llm_params = item.get("parameters")
+                                            if llm_name and isinstance(llm_params, dict):
+                                                parsed_tool_calls_from_content.append(
+                                                    ToolCall(
+                                                        id=item.get(
+                                                            "id",
+                                                            f"call_{uuid.uuid4().hex[:8]}",
+                                                        ),  # Generate an ID
+                                                        type="function",
+                                                        function=ToolCallFunction(
+                                                            name=llm_name,
+                                                            arguments=json.dumps(llm_params),
+                                                        ),
+                                                    )
+                                                )
+                                                continue  # Found valid LLM-specific style tool call
+
+                                    if parsed_tool_calls_from_content:
+                                        logger.debug(
+                                            f"Parsed tool_calls from content field: {parsed_tool_calls_from_content}"
+                                        )
+                                        return GenerationResult(tool_calls=parsed_tool_calls_from_content)
+
+                                    # If JSON but not a recognized tool call, it's just JSON content
+                                    logger.debug(
+                                        "Content was JSON, but not a recognized tool call structure. Treating as text."
+                                    )
+                                    return GenerationResult(content=content_str.strip())
+
+                                except json.JSONDecodeError:
+                                    # Content is not JSON, so it's a regular text response
+                                    logger.debug("Content is not JSON. Treating as text.")
+                                    return GenerationResult(content=content_str.strip())
+
+                        # If neither tool_calls nor parsable content that looks like a tool call
+                        logger.warning(f"Fireworks API response malformed or no actionable content/tool_calls: {data}")
+                        return GenerationResult()
+
+                    # ... (rest of the error handling as before) ...
+                    elif response.status == 429:  # Rate limit
+                        retry_after = int(response.headers.get("Retry-After", "5"))
+                        logger.warning(f"Rate limited. Retrying after {retry_after}s (attempt {attempt+1}).")
+                        await asyncio.sleep(retry_after)
+                    elif response.status in [401, 403]:  # Auth errors
+                        error_text = await response.text()
+                        logger.error(f"Fireworks API Auth Error ({response.status}): {error_text}")
+                        return GenerationResult()  # Empty result on auth error
+                    elif response.status >= 500:  # Server errors
+                        logger.warning(
+                            f"Fireworks API Server Error ({response.status}). Retrying (attempt {attempt+1})."
+                        )
+                        await asyncio.sleep(2**attempt)
+                    else:  # Other client errors
+                        error_text = await response.text()
+                        logger.error(f"Fireworks API request failed ({response.status}): {error_text}")
+                        return GenerationResult()  # Empty result
+            logger.error("Max retries reached for Fireworks API call.")
+            return GenerationResult()
+        except aiohttp.ClientError as e:
+            logger.error(f"AIOHTTP client error: {e}")
+            return GenerationResult()
+        except Exception as e:
+            logger.error(f"Unexpected error in FireworksModelClient: {e}", exc_info=True)
+            return GenerationResult()

eval_protocol/generic_server.py

@@ -0,0 +1,165 @@
+import importlib
+import os
+from typing import Any, Dict, List, Optional
+
+import uvicorn
+from fastapi import Depends, FastAPI, HTTPException, Request
+from pydantic import BaseModel, ValidationError
+
+# Assuming these models are correctly defined in eval_protocol.models
+from eval_protocol.models import EvaluateResult, Message
+
+
+# --- Request and Response Models ---
+class EvaluationRequest(BaseModel):
+    messages: List[Dict[str, Any]]  # Could also be List[Message] if we enforce that model on input
+    ground_truth: Optional[str] = None
+    kwargs: Optional[Dict[str, Any]] = {}
+
+
+# --- Global variable to store the loaded reward function ---
+# This is a simple approach for a single-function server.
+# If multiple functions were to be served by one instance, a different mechanism would be needed.
+_LOADED_REWARD_FUNCTION = None
+_REWARD_FUNCTION_NAME = "N/A"
+
+# --- API Key Authentication Dependency ---
+EXPECTED_API_KEY = os.environ.get("RK_ENDPOINT_API_KEY")
+
+
+async def verify_api_key(request: Request):
+    if EXPECTED_API_KEY:
+        # Check for X-Api-Key header first
+        client_api_key = request.headers.get("X-Api-Key")
+        # If not found, check for Authorization: Bearer <key>
+        if not client_api_key:
+            auth_header = request.headers.get("Authorization")
+            if auth_header and auth_header.startswith("Bearer "):
+                client_api_key = auth_header.split(" ", 1)[1]
+
+        if not client_api_key:
+            raise HTTPException(status_code=401, detail="API key required but not provided.")
+        if client_api_key != EXPECTED_API_KEY:
+            raise HTTPException(status_code=403, detail="Invalid API key.")
+    return True  # Allow request if no key expected or if key is valid
+
+
+# --- FastAPI App ---
+app = FastAPI(
+    title="Reward Kit Generic Reward Function Server",
+    description="Serves a dynamically loaded reward function.",
+    version="0.1.0",  # Or use eval_protocol.__version__
+)
+
+
+@app.post("/evaluate", response_model=EvaluateResult, dependencies=[Depends(verify_api_key)])
+async def evaluate_endpoint(request: EvaluationRequest):
+    """
+    Endpoint to evaluate a given set of messages using the loaded reward function.
+    Requires API key if RK_ENDPOINT_API_KEY environment variable is set.
+    """
+    if _LOADED_REWARD_FUNCTION is None:
+        raise HTTPException(status_code=500, detail="Reward function not loaded.")
+
+    try:
+        # The user's reward function is expected to match the @reward_function signature
+        func_args = {
+            "messages": request.messages,
+            "ground_truth": request.ground_truth,
+            **(request.kwargs or {}),
+        }
+
+        result = _LOADED_REWARD_FUNCTION(**func_args)
+
+        if not isinstance(result, EvaluateResult):
+            # This case should ideally not happen if functions are correctly decorated
+            # and return EvaluateResult, but good to have a fallback.
+            print(
+                f"Warning: Reward function '{_REWARD_FUNCTION_NAME}' did not return an EvaluateResult instance. Type: {type(result)}"
+            )
+            # Attempt to construct an EvaluateResult if it's a dict-like object,
+            # otherwise, this will raise an error or return a poorly formed response.
+            # For robustness, one might want to wrap this in another try-except.
+            return EvaluateResult(
+                score=0.0,
+                reason="Invalid return type from reward function, check server logs.",
+                is_score_valid=False,
+                metrics={},
+            )
+
+        return result
+    except ValidationError as ve:  # Pydantic validation error from reward function's input/output
+        print(f"Validation Error calling reward function '{_REWARD_FUNCTION_NAME}': {ve}")
+        raise HTTPException(
+            status_code=422,
+            detail=f"Input/Output validation error for reward function: {ve.errors()}",
+        )
+    except Exception as e:
+        print(f"Error during evaluation with reward function '{_REWARD_FUNCTION_NAME}': {e}")
+        # Consider logging the full traceback here
+        raise HTTPException(status_code=500, detail=f"Internal server error during evaluation: {str(e)}")
+
+
+@app.get("/health")
+async def health_check():
+    """
+    Health check endpoint.
+    """
+    if _LOADED_REWARD_FUNCTION:
+        return {"status": "ok", "reward_function": _REWARD_FUNCTION_NAME}
+    else:
+        return {"status": "error", "reason": "Reward function not loaded"}
+
+
+def load_reward_function(import_string: str):
+    """
+    Loads a reward function from an import string (e.g., 'my_module.my_function').
+    """
+    global _LOADED_REWARD_FUNCTION, _REWARD_FUNCTION_NAME
+    try:
+        module_path, function_name = import_string.rsplit(".", 1)
+        module = importlib.import_module(module_path)
+        _LOADED_REWARD_FUNCTION = getattr(module, function_name)
+        _REWARD_FUNCTION_NAME = import_string
+        print(f"Successfully loaded reward function: {_REWARD_FUNCTION_NAME}")
+    except Exception as e:
+        print(f"Error loading reward function from '{import_string}': {e}")
+        _LOADED_REWARD_FUNCTION = None
+        _REWARD_FUNCTION_NAME = "Error loading"
+        raise  # Re-raise to make it fatal if loading fails on startup
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Run the Generic Reward Function Server.")
+    parser.add_argument(
+        "import_string",
+        type=str,
+        help="Import string for the reward function (e.g., 'my_package.my_module.my_reward_function')",
+    )
+    parser.add_argument("--host", type=str, default="127.0.0.1", help="Host to bind the server to.")
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=8080,  # Standard port for Cloud Run, etc.
+        help="Port to bind the server to.",
+    )
+    # Add --reload for uvicorn if needed for development
+    # parser.add_argument("--reload", action="store_true", help="Enable auto-reload for development.")
+
+    args = parser.parse_args()
+
+    try:
+        load_reward_function(args.import_string)
+    except Exception:
+        print(f"Failed to load reward function. Exiting.")
+        exit(1)
+
+    if not _LOADED_REWARD_FUNCTION:
+        print(f"Reward function {_REWARD_FUNCTION_NAME} could not be loaded. Server will not start correctly.")
+        # Depending on desired behavior, could exit here or let it run and fail on /evaluate
+        exit(1)
+
+    print(f"Starting server for reward function: {args.import_string} on http://{args.host}:{args.port}")
+    uvicorn.run(app, host=args.host, port=args.port)  # reload=args.reload for dev

eval_protocol/integrations/__init__.py

@@ -0,0 +1,12 @@
+"""Integration helpers for Reward Kit."""
+
+from .braintrust import reward_fn_to_scorer, scorer_to_reward_fn
+from .openeval import adapt
+from .trl import create_trl_adapter
+
+__all__ = [
+    "adapt",
+    "scorer_to_reward_fn",
+    "reward_fn_to_scorer",
+    "create_trl_adapter",
+]

eval_protocol/integrations/braintrust.py

@@ -0,0 +1,51 @@
+"""Adapters for integrating Reward Kit with Braintrust scoring functions."""
+
+from typing import Any, Callable, List, Optional
+
+from eval_protocol.models import EvaluateResult, Message
+from eval_protocol.typed_interface import reward_function
+
+# Type alias for Braintrust scoring functions
+BraintrustScorer = Callable[[Any, Any, Any], float]
+
+
+def scorer_to_reward_fn(
+    scorer: BraintrustScorer,
+    *,
+    messages_to_input: Optional[Callable[[List[Message]], Any]] = None,
+    ground_truth_to_expected: Optional[Callable[[List[Message]], Any]] = None,
+) -> Callable[[List[Message], Optional[List[Message]]], EvaluateResult]:
+    """Wrap a Braintrust scorer as a Reward Kit reward function."""
+
+    @reward_function
+    def reward_fn(messages: List[Message], ground_truth: Optional[List[Message]] = None, **kwargs) -> EvaluateResult:
+        input_val = messages_to_input(messages) if messages_to_input else messages[0].content
+        output_val = messages[-1].content
+        expected_val = None
+        if ground_truth:
+            expected_val = (
+                ground_truth_to_expected(ground_truth) if ground_truth_to_expected else ground_truth[-1].content
+            )
+        score = scorer(input_val, output_val, expected_val)
+        return EvaluateResult(score=score)
+
+    return reward_fn
+
+
+def reward_fn_to_scorer(
+    reward_fn: Callable[[List[Message], Optional[List[Message]]], EvaluateResult],
+) -> BraintrustScorer:
+    """Create a Braintrust-compatible scorer from a Reward Kit reward function."""
+
+    def scorer(input_val: Any, output: Any, expected: Any) -> float:
+        messages = [
+            Message(role="user", content=str(input_val)),
+            Message(role="assistant", content=str(output)),
+        ]
+        ground_truth = None
+        if expected is not None:
+            ground_truth = [Message(role="assistant", content=str(expected))]
+        result = reward_fn(messages=messages, ground_truth=ground_truth)
+        return result.score
+
+    return scorer

eval_protocol/integrations/deepeval.py

@@ -0,0 +1,106 @@
+from typing import Any, Dict, List, Optional
+
+from eval_protocol.models import EvaluateResult, MetricResult
+from eval_protocol.typed_interface import reward_function
+
+__all__ = ["adapt_metric"]
+
+try:
+    from deepeval.metrics.base_metric import BaseConversationalMetric, BaseMetric
+    from deepeval.test_case import ConversationalTestCase, LLMTestCase
+except Exception:  # pragma: no cover - deepeval is optional
+    BaseMetric = None
+    BaseConversationalMetric = None
+    LLMTestCase = None
+    ConversationalTestCase = None
+
+
+def _metric_name(metric: Any) -> str:
+    name = getattr(metric, "__name__", None)
+    if name and name not in {
+        "Base Metric",
+        "Base Conversational Metric",
+        "Base Multimodal Metric",
+    }:
+        return str(name)
+    name = getattr(metric, "name", None)
+    if name:
+        return str(name)
+    return metric.__class__.__name__
+
+
+def adapt_metric(metric: Any):
+    """Adapt a deepeval metric object into a reward-kit reward function."""
+
+    @reward_function
+    def wrapped(
+        messages: List[Dict[str, Any]],
+        ground_truth: Optional[str] = None,
+        **kwargs: Any,
+    ) -> EvaluateResult:
+        if BaseMetric is None or LLMTestCase is None:
+            raise ImportError("deepeval must be installed to use this integration")
+        if not messages:
+            return EvaluateResult(score=0.0, reason="No messages", metrics={})
+
+        output = messages[-1].get("content", "")
+        input_msg = ""
+        if len(messages) >= 2:
+            input_msg = messages[-2].get("content", "")
+
+        def _build_case_kwargs() -> Dict[str, Any]:
+            case_kwargs: Dict[str, Any] = {}
+            params = getattr(metric, "evaluation_params", None)
+            if params:
+                for param in params:
+                    if param.value == "input":
+                        case_kwargs["input"] = input_msg
+                    elif param.value == "actual_output":
+                        case_kwargs["actual_output"] = output
+                    elif param.value == "expected_output":
+                        case_kwargs["expected_output"] = ground_truth
+                    elif param.value == "context":
+                        case_kwargs["context"] = kwargs.get("context")
+                    elif param.value == "retrieval_context":
+                        case_kwargs["retrieval_context"] = kwargs.get("retrieval_context")
+                    elif param.value == "tools_called":
+                        case_kwargs["tools_called"] = kwargs.get("tools_called")
+                    elif param.value == "expected_tools":
+                        case_kwargs["expected_tools"] = kwargs.get("expected_tools")
+            else:
+                case_kwargs = {
+                    "input": input_msg,
+                    "actual_output": output,
+                    "expected_output": ground_truth,
+                }
+            if "input" not in case_kwargs:
+                case_kwargs["input"] = input_msg
+            if "actual_output" not in case_kwargs:
+                case_kwargs["actual_output"] = output
+            return case_kwargs
+
+        if isinstance(metric, BaseConversationalMetric):
+            turns = []
+            for i, msg in enumerate(messages):
+                turn_input = messages[i - 1].get("content", "") if i > 0 else ""
+                output_turn = msg.get("content", "")
+                input_msg_backup = input_msg
+                input_msg = turn_input
+                output = output_turn
+                turn_kwargs = _build_case_kwargs()
+                turns.append(LLMTestCase(**turn_kwargs))
+                input_msg = input_msg_backup
+                output = messages[-1].get("content", "")
+            test_case = ConversationalTestCase(turns=turns)
+        else:
+            case_kwargs = _build_case_kwargs()
+            test_case = LLMTestCase(**case_kwargs)
+
+        metric.measure(test_case, **kwargs)
+        score = float(metric.score or 0.0)
+        reason = getattr(metric, "reason", None)
+        name = _metric_name(metric)
+        metrics = {name: MetricResult(score=score, reason=reason or "", is_score_valid=True)}
+        return EvaluateResult(score=score, reason=reason, metrics=metrics)
+
+    return wrapped

eval_protocol/integrations/openeval.py

@@ -0,0 +1,40 @@
+from typing import Any, Callable, Dict, List, Union
+
+from eval_protocol.models import EvaluateResult, MetricResult
+from eval_protocol.typed_interface import reward_function
+
+__all__ = ["adapt"]
+
+
+def _convert_result(res: Dict[str, Any]) -> EvaluateResult:
+    score = float(res.get("score", 0.0))
+    reason = res.get("comment")
+    key = res.get("key", "openeval")
+    metrics = {key: MetricResult(score=score, reason=reason or "", is_score_valid=True)}
+    return EvaluateResult(score=score, reason=reason, metrics=metrics)
+
+
+def adapt(openeval_fn: Callable[..., Union[Dict[str, Any], List[Dict[str, Any]]]]):
+    """Adapt an OpenEvals evaluator into a reward-kit reward function."""
+
+    @reward_function
+    def wrapped(
+        messages: List[Dict[str, str]],
+        ground_truth: Union[str, List[Dict[str, str]], None] = None,
+        **kwargs: Any,
+    ) -> EvaluateResult:
+        if not messages:
+            return EvaluateResult(score=0.0, reason="No messages", metrics={})
+        output = messages[-1].get("content", "")
+        reference = None
+        if isinstance(ground_truth, list):
+            if ground_truth:
+                reference = ground_truth[-1].get("content")
+        else:
+            reference = ground_truth
+        res = openeval_fn(outputs=output, reference_outputs=reference, **kwargs)
+        if isinstance(res, list):
+            res = res[0]
+        return _convert_result(res)
+
+    return wrapped