agent-learning-kit 0.1.0__py3-none-any.whl

fi/simulate/agent/wrapper.py

@@ -0,0 +1,99 @@
+from typing import List, Dict, Union, Any, Optional, Literal
+from pydantic import BaseModel, Field
+from abc import ABC, abstractmethod
+
+
+ArtifactType = Literal[
+    "text",
+    "image",
+    "audio",
+    "video",
+    "screenshot",
+    "browser_dom",
+    "file",
+    "json",
+    "trace",
+]
+
+
+class SimulationArtifact(BaseModel):
+    """
+    Modality-neutral artifact carried through a simulation.
+
+    Use `uri` or `path` for large media, `data` for small inline payloads, and
+    `metadata` for framework-specific details like sample rate, viewport, page
+    URL, or image dimensions.
+    """
+
+    type: ArtifactType
+    uri: Optional[str] = None
+    path: Optional[str] = None
+    data: Optional[Any] = None
+    mime_type: Optional[str] = None
+    role: Optional[str] = None
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+
+
+class SimulationEvent(BaseModel):
+    """Normalized event for tools, memory, browser/CUA actions, voice states, and framework spans."""
+
+    type: str
+    name: Optional[str] = None
+    payload: Dict[str, Any] = Field(default_factory=dict)
+    timestamp_ms: Optional[int] = None
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+
+
+class AgentInput(BaseModel):
+    """
+    Input data passed to the user's agent wrapper during a simulation step.
+    """
+    thread_id: str
+    messages: List[Dict[str, Any]]  # Full conversation history: [{"role": "user", "content": "..."}]
+    new_message: Optional[Dict[str, Any]] = None # The latest message to respond to
+    
+    # Metadata for execution context (useful for logging/debugging)
+    execution_id: Optional[str] = None
+    turn_index: Optional[int] = None
+    scenario_name: Optional[str] = None
+    persona: Optional[Dict[str, Any]] = None
+    situation: Optional[str] = None
+    expected_outcome: Optional[str] = None
+    modality: Optional[str] = None
+    artifacts: List[SimulationArtifact] = Field(default_factory=list)
+    events: List[SimulationEvent] = Field(default_factory=list)
+    memory: Dict[str, Any] = Field(default_factory=dict)
+    tools: List[Dict[str, Any]] = Field(default_factory=list)
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+
+class AgentResponse(BaseModel):
+    """
+    Standardized response from the user's agent.
+    """
+    content: str
+    tool_calls: Optional[List[Dict[str, Any]]] = None
+    tool_responses: Optional[List[Dict[str, Any]]] = None  # Tool role messages with results
+    artifacts: List[SimulationArtifact] = Field(default_factory=list)
+    events: List[SimulationEvent] = Field(default_factory=list)
+    memory_updates: Optional[Dict[str, Any]] = None
+    state: Optional[Dict[str, Any]] = None
+    metadata: Optional[Dict[str, Any]] = None
+
+class AgentWrapper(ABC):
+    """
+    Base class for wrapping user agents to work with the simulation SDK.
+    Users should implement the `call` method.
+    """
+    
+    @abstractmethod
+    async def call(self, input: AgentInput) -> Union[str, AgentResponse]:
+        """
+        Process the input and return the agent's response.
+        
+        Args:
+            input: The AgentInput object containing message history and context.
+            
+        Returns:
+            A string (content only) or AgentResponse object.
+        """
+        pass

fi/simulate/agent/wrappers/__init__.py

@@ -0,0 +1,18 @@
+from fi.simulate.agent.wrappers.openai import OpenAIAgentWrapper
+from fi.simulate.agent.wrappers.langchain import LangChainAgentWrapper
+from fi.simulate.agent.wrappers.gemini import GeminiAgentWrapper
+from fi.simulate.agent.wrappers.anthropic import AnthropicAgentWrapper
+from fi.simulate.agent.wrappers.http import HTTPAgentWrapper
+from fi.simulate.agent.wrappers.websocket import WebSocketAgentWrapper
+
+OpenAICompatibleHTTPAgentWrapper = HTTPAgentWrapper
+
+__all__ = [
+    "OpenAIAgentWrapper",
+    "LangChainAgentWrapper",
+    "GeminiAgentWrapper",
+    "AnthropicAgentWrapper",
+    "HTTPAgentWrapper",
+    "OpenAICompatibleHTTPAgentWrapper",
+    "WebSocketAgentWrapper",
+]

fi/simulate/agent/wrappers/anthropic.py

@@ -0,0 +1,62 @@
+from typing import Any, Union
+from fi.simulate.agent.wrapper import AgentWrapper, AgentInput, AgentResponse
+
+class AnthropicAgentWrapper(AgentWrapper):
+    """
+    Wrapper for Anthropic (Claude) agents.
+    Automatically handles message conversion to Anthropic format.
+    """
+    def __init__(self, client: Any, model: str = "claude-sonnet-4-5-20250929", system_prompt: str = None, max_tokens: int = 1024):
+        """
+        Args:
+            client: The Anthropic client instance (AsyncAnthropic or Anthropic).
+            model: The model name to use.
+            system_prompt: Optional system instructions for the agent.
+            max_tokens: Maximum number of tokens to generate (default: 1024).
+        """
+        self.client = client
+        self.model = model
+        self.system_prompt = system_prompt
+        self.max_tokens = max_tokens
+
+    async def call(self, input: AgentInput) -> Union[str, AgentResponse]:
+        # Convert internal message format to Anthropic format
+        # Anthropic messages API expects: [{"role": "user"|"assistant", "content": "..."}]
+        # It does NOT support "system" role in the messages list; system prompt is a top-level param.
+        
+        messages = []
+        # Use configured system prompt by default
+        system_prompt = self.system_prompt
+        
+        for msg in input.messages:
+            if msg["role"] == "system":
+                # If history has system message (unlikely due to filtering), it overrides? 
+                # Or we ignore it to respect wrapper config? 
+                # Let's check if it exists and use it if self.system_prompt is None
+                if system_prompt is None:
+                    system_prompt = msg["content"]
+            else:
+                messages.append({
+                    "role": msg["role"],
+                    "content": msg["content"]
+                })
+        
+        # Check for AsyncAnthropic vs Sync
+        # Heuristic: check for 'messages.create' and if client class name contains Async
+        is_async = type(self.client).__name__.startswith("Async")
+        
+        kwargs = {
+            "model": self.model,
+            "max_tokens": self.max_tokens,
+            "messages": messages
+        }
+        if system_prompt:
+            kwargs["system"] = system_prompt
+
+        if is_async:
+            message = await self.client.messages.create(**kwargs)
+        else:
+            message = self.client.messages.create(**kwargs)
+
+        return message.content[0].text
+

fi/simulate/agent/wrappers/gemini.py

@@ -0,0 +1,65 @@
+from typing import Any, Union
+from fi.simulate.agent.wrapper import AgentWrapper, AgentInput, AgentResponse
+
+class GeminiAgentWrapper(AgentWrapper):
+    """
+    Wrapper for Google Gemini (Generative AI) agents.
+    Supports google-generativeai SDK.
+    """
+    def __init__(self, model: Any, system_prompt: str = None):
+        """
+        Args:
+            model: An instance of google.generativeai.GenerativeModel
+            system_prompt: Optional system instructions. 
+                           Note: Ideally configure system_instruction on the model itself.
+                           If provided here, it will be prepended as a user message.
+        """
+        self.model = model
+        self.system_prompt = system_prompt
+
+    async def call(self, input: AgentInput) -> Union[str, AgentResponse]:
+        # Convert internal messages to Gemini format (Content objects)
+        # Note: Gemini SDK manages chat history via ChatSession usually,
+        # but for stateless call we pass full history if supported, 
+        # or we might need to reconstruct a chat session.
+        
+        # Simple reconstruction of history for a chat session
+        history = []
+        
+        if self.system_prompt:
+            # Prepend system prompt as a user message for context
+            history.append({"role": "user", "parts": [f"System Instruction: {self.system_prompt}"]})
+            # Add a dummy model acknowledgement to keep turns valid (User -> Model -> User)
+            history.append({"role": "model", "parts": ["Understood."]})
+            
+        for msg in input.messages:
+            role = "user" if msg["role"] == "user" else "model"
+            content = msg["content"]
+            
+            # Gemini typically expects history excluding the last message which is passed to send_message
+            history.append({"role": role, "parts": [content]})
+            
+        if not history:
+            raise ValueError("No messages provided to Gemini wrapper")
+
+        # The last user message is the prompt
+        last_turn = history.pop()
+        if last_turn["role"] != "user":
+            # If the last message wasn't user, something is weird in the flow,
+            # but we can try to send empty or handle it. 
+            # Ideally simulator sends User message last.
+            prompt = ""
+        else:
+            prompt = last_turn["parts"][0]
+
+        # Start a chat with the history
+        chat = self.model.start_chat(history=history)
+        
+        # Check if async generation is supported (google-generativeai >= 0.3.0 has send_message_async)
+        if hasattr(chat, "send_message_async"):
+            response = await chat.send_message_async(prompt)
+        else:
+            # Fallback to sync
+            response = chat.send_message(prompt)
+            
+        return response.text

fi/simulate/agent/wrappers/http.py

@@ -0,0 +1,404 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import time
+import urllib.error
+import urllib.request
+from typing import Any, Mapping, Optional, Sequence
+from urllib.parse import urlparse
+
+from fi.simulate.agent.wrapper import (
+    AgentInput,
+    AgentResponse,
+    SimulationArtifact,
+    SimulationEvent,
+)
+from fi.simulate.agent.wrapper import AgentWrapper
+
+
+class HTTPAgentWrapper(AgentWrapper):
+    """HTTP/OpenAI-compatible target adapter for external agent simulation."""
+
+    def __init__(
+        self,
+        *,
+        endpoint: str,
+        protocol: str = "fi.alk",
+        model: Optional[str] = None,
+        api_key: Optional[str] = None,
+        api_key_env: Optional[str] = None,
+        headers: Optional[Mapping[str, str]] = None,
+        timeout: float = 30.0,
+        include_tools: bool = True,
+        system_prompt: Optional[str] = None,
+        metadata: Optional[Mapping[str, Any]] = None,
+    ) -> None:
+        if not endpoint:
+            raise ValueError("endpoint is required")
+        self.endpoint = endpoint
+        self.protocol = _normalize_protocol(protocol)
+        self.model = model
+        self.api_key = api_key
+        self.api_key_env = api_key_env
+        self.headers = {str(k): str(v) for k, v in dict(headers or {}).items()}
+        self.timeout = float(timeout)
+        self.include_tools = bool(include_tools)
+        self.system_prompt = system_prompt
+        self.metadata = dict(metadata or {})
+
+    async def call(self, input: AgentInput) -> AgentResponse:
+        started = time.time()
+        request_payload = self._request_payload(input)
+        headers = self._request_headers()
+        status_code = 0
+        response_payload: dict[str, Any] = {}
+        error: Optional[str] = None
+        try:
+            status_code, response_payload = await asyncio.to_thread(
+                self._post_json,
+                request_payload,
+                headers,
+            )
+            if status_code >= 400:
+                error = _response_error_text(response_payload) or (
+                    f"HTTP target returned status {status_code}"
+                )
+            response = self._agent_response_from_payload(response_payload)
+        except Exception as exc:
+            error = str(exc)
+            response = AgentResponse(content=f"HTTP target failed: {exc}")
+
+        latency_ms = round((time.time() - started) * 1000, 4)
+        trace = {
+            "kind": "external_agent_http_trace",
+            "protocol": self.protocol,
+            "endpoint": _redacted_endpoint(self.endpoint),
+            "endpoint_host": urlparse(self.endpoint).netloc,
+            "model": self.model,
+            "status_code": status_code,
+            "latency_ms": latency_ms,
+            "request_message_count": len(input.messages),
+            "request_tool_count": len(input.tools) if self.include_tools else 0,
+            "response_tool_call_count": len(response.tool_calls or []),
+            "success": error is None and 200 <= status_code < 300,
+            "request_header_names": sorted(headers),
+            "auth": {
+                "mode": "bearer" if self._resolved_api_key() else "none",
+                "api_key_env": self.api_key_env,
+                "redacted": bool(self._resolved_api_key()),
+            },
+            "error": error,
+            **self.metadata,
+        }
+        response.events.append(
+            SimulationEvent(
+                type="external_agent",
+                name="external_agent_http_call",
+                payload=trace,
+            )
+        )
+        response.artifacts.append(
+            SimulationArtifact(
+                type="trace",
+                role="agent",
+                data=trace,
+                metadata={"kind": "external_agent_http_trace"},
+            )
+        )
+        state = dict(response.state or {})
+        state["external_agent"] = trace
+        state["external_agent_trace"] = trace
+        response.state = state
+        metadata = dict(response.metadata or {})
+        metadata["external_agent"] = trace
+        metadata["external_agent_trace"] = trace
+        response.metadata = metadata
+        return response
+
+    def _request_payload(self, input: AgentInput) -> dict[str, Any]:
+        messages = _messages_for_protocol(input.messages, self.protocol)
+        if self.system_prompt:
+            messages = [{"role": "system", "content": self.system_prompt}, *messages]
+        if self.protocol == "openai_chat":
+            payload: dict[str, Any] = {
+                "model": self.model or "agent-learning-target",
+                "messages": messages,
+            }
+            if self.include_tools and input.tools:
+                payload["tools"] = [_openai_tool_spec(tool) for tool in input.tools]
+                payload["tool_choice"] = "auto"
+            return payload
+        return {
+            "thread_id": input.thread_id,
+            "execution_id": input.execution_id,
+            "turn_index": input.turn_index,
+            "scenario_name": input.scenario_name,
+            "persona": input.persona,
+            "situation": input.situation,
+            "expected_outcome": input.expected_outcome,
+            "messages": messages,
+            "new_message": input.new_message,
+            "tools": list(input.tools) if self.include_tools else [],
+            "metadata": input.metadata,
+        }
+
+    def _request_headers(self) -> dict[str, str]:
+        headers = {"Content-Type": "application/json", **self.headers}
+        api_key = self._resolved_api_key()
+        if api_key and not any(key.lower() == "authorization" for key in headers):
+            headers["Authorization"] = f"Bearer {api_key}"
+        return headers
+
+    def _resolved_api_key(self) -> str:
+        if self.api_key not in (None, ""):
+            return str(self.api_key)
+        if self.api_key_env:
+            return os.environ.get(self.api_key_env, "")
+        return ""
+
+    def _post_json(
+        self,
+        payload: Mapping[str, Any],
+        headers: Mapping[str, str],
+    ) -> tuple[int, dict[str, Any]]:
+        body = json.dumps(payload, default=str).encode("utf-8")
+        request = urllib.request.Request(
+            self.endpoint,
+            data=body,
+            headers=dict(headers),
+            method="POST",
+        )
+        try:
+            with urllib.request.urlopen(request, timeout=self.timeout) as response:
+                status = int(getattr(response, "status", 200))
+                text = response.read().decode("utf-8")
+        except urllib.error.HTTPError as exc:
+            status = int(exc.code)
+            text = exc.read().decode("utf-8")
+        if not text:
+            return status, {}
+        try:
+            parsed = json.loads(text)
+        except json.JSONDecodeError as exc:
+            raise ValueError(f"HTTP target returned non-JSON response: {exc}") from exc
+        if not isinstance(parsed, dict):
+            raise ValueError("HTTP target response must be a JSON object")
+        return status, parsed
+
+    def _agent_response_from_payload(self, payload: Mapping[str, Any]) -> AgentResponse:
+        if self.protocol == "openai_chat":
+            message = _openai_message(payload)
+            return AgentResponse(
+                content=_content_text(message.get("content")),
+                tool_calls=_openai_tool_calls(message.get("tool_calls")),
+                metadata={
+                    "finish_reason": _openai_finish_reason(payload),
+                    "usage": dict(payload.get("usage") or {}),
+                },
+            )
+        return AgentResponse(
+            content=_content_text(payload.get("content") or payload.get("message")),
+            tool_calls=_tool_call_list(payload.get("tool_calls")),
+            tool_responses=_tool_response_list(payload.get("tool_responses")),
+            artifacts=_artifact_list(payload.get("artifacts")),
+            events=_event_list(payload.get("events")),
+            memory_updates=_optional_mapping(payload.get("memory_updates")),
+            state=_optional_mapping(payload.get("state")),
+            metadata=_optional_mapping(payload.get("metadata")),
+        )
+
+
+def _normalize_protocol(value: str) -> str:
+    protocol = str(value or "fi.alk").lower().replace("-", "_")
+    aliases = {
+        "openai": "openai_chat",
+        "openai_compatible": "openai_chat",
+        "chat_completions": "openai_chat",
+        "agent_learning_http": "fi.alk",
+        "http": "fi.alk",
+    }
+    protocol = aliases.get(protocol, protocol)
+    if protocol not in {"fi.alk", "openai_chat"}:
+        raise ValueError("protocol must be one of: fi.alk, openai_chat")
+    return protocol
+
+
+def _messages_for_protocol(
+    messages: Sequence[Mapping[str, Any]], protocol: str
+) -> list[dict[str, Any]]:
+    """Encode canonical ALK history for the selected external-agent protocol.
+
+    ALK keeps tool calls structured internally. OpenAI-compatible endpoints require that array
+    unchanged, while the FutureAGI callback ``AgentInput`` schema represents historical
+    ``tool_calls`` as a JSON string. Normalizing here keeps the conversation runner generic and
+    prevents a retry from accidentally executing a side-effecting tool twice.
+    """
+    normalized = [dict(message) for message in messages]
+    if protocol != "fi.alk":
+        return normalized
+    for message in normalized:
+        tool_calls = message.get("tool_calls")
+        if isinstance(tool_calls, Sequence) and not isinstance(
+            tool_calls, (str, bytes)
+        ):
+            message["tool_calls"] = json.dumps(
+                list(tool_calls), separators=(",", ":"), default=str
+            )
+    return normalized
+
+
+def _openai_tool_spec(tool: Mapping[str, Any]) -> dict[str, Any]:
+    # Accept both the flat SDK tool shape ({name, description, parameters}) and
+    # the OpenAI-nested shape ({"type": "function", "function": {...}}). Without
+    # reading the nested ``function`` block, a nested spec loses its name and the
+    # model is handed a tool literally called "tool" — so it can never call the
+    # real tool and the environment's mock never matches.
+    fn = tool.get("function") if isinstance(tool.get("function"), Mapping) else {}
+    name = str(
+        tool.get("name")
+        or fn.get("name")
+        or tool.get("tool")
+        or tool.get("id")
+        or "tool"
+    )
+    parameters = tool.get("parameters")
+    if not isinstance(parameters, Mapping):
+        parameters = fn.get("parameters")
+    if not isinstance(parameters, Mapping):
+        parameters = {"type": "object", "properties": {}}
+    description = tool.get("description") or fn.get("description") or f"Tool {name}"
+    return {
+        "type": "function",
+        "function": {
+            "name": name,
+            "description": str(description),
+            "parameters": dict(parameters),
+        },
+    }
+
+
+def _openai_message(payload: Mapping[str, Any]) -> dict[str, Any]:
+    choices = payload.get("choices")
+    if isinstance(choices, Sequence) and not isinstance(choices, (str, bytes)):
+        if choices:
+            choice = choices[0]
+            if isinstance(choice, Mapping):
+                message = choice.get("message")
+                if isinstance(message, Mapping):
+                    return dict(message)
+    message = payload.get("message")
+    return dict(message) if isinstance(message, Mapping) else dict(payload)
+
+
+def _openai_finish_reason(payload: Mapping[str, Any]) -> Optional[str]:
+    choices = payload.get("choices")
+    if isinstance(choices, Sequence) and not isinstance(choices, (str, bytes)):
+        if choices and isinstance(choices[0], Mapping):
+            value = choices[0].get("finish_reason")
+            return str(value) if value is not None else None
+    return None
+
+
+def _openai_tool_calls(value: Any) -> list[dict[str, Any]]:
+    calls = _tool_call_list(value)
+    normalized: list[dict[str, Any]] = []
+    for index, call in enumerate(calls, start=1):
+        function = call.get("function")
+        if isinstance(function, Mapping):
+            name = function.get("name")
+            arguments = function.get("arguments", {})
+        else:
+            name = call.get("name") or call.get("tool")
+            arguments = call.get("arguments", call.get("args", {}))
+        normalized.append(
+            {
+                "id": str(call.get("id") or f"call_{index}"),
+                "type": str(call.get("type") or "function"),
+                "function": {
+                    "name": str(name or ""),
+                    "arguments": (
+                        arguments
+                        if isinstance(arguments, str)
+                        else json.dumps(arguments or {}, default=str)
+                    ),
+                },
+            }
+        )
+    return normalized
+
+
+def _tool_call_list(value: Any) -> list[dict[str, Any]]:
+    if not isinstance(value, Sequence) or isinstance(value, (str, bytes)):
+        return []
+    return [dict(item) for item in value if isinstance(item, Mapping)]
+
+
+def _tool_response_list(value: Any) -> list[dict[str, Any]]:
+    if not isinstance(value, Sequence) or isinstance(value, (str, bytes)):
+        return []
+    return [dict(item) for item in value if isinstance(item, Mapping)]
+
+
+def _artifact_list(value: Any) -> list[SimulationArtifact]:
+    if not isinstance(value, Sequence) or isinstance(value, (str, bytes)):
+        return []
+    artifacts: list[SimulationArtifact] = []
+    for item in value:
+        if isinstance(item, SimulationArtifact):
+            artifacts.append(item)
+        elif isinstance(item, Mapping):
+            artifacts.append(SimulationArtifact(**dict(item)))
+    return artifacts
+
+
+def _event_list(value: Any) -> list[SimulationEvent]:
+    if not isinstance(value, Sequence) or isinstance(value, (str, bytes)):
+        return []
+    events: list[SimulationEvent] = []
+    for item in value:
+        if isinstance(item, SimulationEvent):
+            events.append(item)
+        elif isinstance(item, Mapping):
+            events.append(SimulationEvent(**dict(item)))
+    return events
+
+
+def _optional_mapping(value: Any) -> Optional[dict[str, Any]]:
+    return dict(value) if isinstance(value, Mapping) else None
+
+
+def _content_text(value: Any) -> str:
+    if isinstance(value, str):
+        return value
+    if isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
+        parts: list[str] = []
+        for item in value:
+            if isinstance(item, Mapping):
+                text = item.get("text") or item.get("content") or item.get("refusal")
+                if text not in (None, ""):
+                    parts.append(str(text))
+            elif item not in (None, ""):
+                parts.append(str(item))
+        return "\n".join(parts)
+    return "" if value is None else str(value)
+
+
+def _response_error_text(payload: Mapping[str, Any]) -> str:
+    error = payload.get("error")
+    if isinstance(error, Mapping):
+        return _content_text(error.get("message") or error.get("detail") or error)
+    if error not in (None, ""):
+        return _content_text(error)
+    for key in ("detail", "message", "status"):
+        if payload.get(key) not in (None, ""):
+            return _content_text(payload.get(key))
+    return ""
+
+
+def _redacted_endpoint(endpoint: str) -> str:
+    parsed = urlparse(endpoint)
+    if not parsed.query:
+        return endpoint
+    return parsed._replace(query="<redacted>").geturl()