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

fi/simulate/agent/wrappers/langchain.py

@@ -0,0 +1,80 @@
+from typing import Any, Union
+from fi.simulate.agent.wrapper import AgentWrapper, AgentInput, AgentResponse
+
+try:
+    from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
+except ImportError:
+    # LangChain is optional. Fallback message classes keep wrappers and tests
+    # usable for mocked/local simulations without installing langchain-core.
+    class _FallbackMessage:
+        def __init__(self, content: str):
+            self.content = content
+
+    class HumanMessage(_FallbackMessage):
+        pass
+
+    class AIMessage(_FallbackMessage):
+        pass
+
+    class SystemMessage(_FallbackMessage):
+        pass
+
+    LANGCHAIN_INSTALLED = False
+else:
+    LANGCHAIN_INSTALLED = True
+
+class LangChainAgentWrapper(AgentWrapper):
+    """
+    Wrapper for LangChain Runnable or Chain agents.
+    """
+    def __init__(self, agent: Any, system_prompt: str = None):
+        """
+        Args:
+            agent: A LangChain Runnable (chain, agent executor) that accepts input.
+                   It is expected to accept a dictionary with "messages" or "input".
+            system_prompt: Optional system prompt to prepend to message history.
+        """
+        self.agent = agent
+        self.system_prompt = system_prompt
+
+    async def call(self, input: AgentInput) -> Union[str, AgentResponse]:
+
+        # Convert history to LangChain messages
+        lc_messages = []
+        
+        if self.system_prompt:
+            lc_messages.append(SystemMessage(content=self.system_prompt))
+            
+        for msg in input.messages:
+            role = msg.get("role")
+            content = msg.get("content")
+            if role == "user":
+                lc_messages.append(HumanMessage(content=content))
+            elif role == "assistant":
+                lc_messages.append(AIMessage(content=content))
+            elif role == "system":
+                lc_messages.append(SystemMessage(content=content))
+        
+        # Invoke the agent
+        # We try standard keys used in LC agents
+        inputs = {
+            "messages": lc_messages,
+            "input": input.new_message.get("content") if input.new_message else "",
+            "chat_history": lc_messages[:-1] if lc_messages else []
+        }
+        
+        # Support both ainvoke and invoke
+        if hasattr(self.agent, "ainvoke"):
+            response = await self.agent.ainvoke(inputs)
+        else:
+            response = self.agent.invoke(inputs)
+            
+        # Parse response
+        if isinstance(response, str):
+            return response
+        elif hasattr(response, "content"):
+            return response.content
+        elif isinstance(response, dict) and "output" in response:
+            return response["output"]
+            
+        return str(response)

fi/simulate/agent/wrappers/openai.py

@@ -0,0 +1,75 @@
+from typing import Any, Union
+from fi.simulate.agent.wrapper import AgentWrapper, AgentInput, AgentResponse
+
+class OpenAIAgentWrapper(AgentWrapper):
+    """
+    Wrapper for OpenAI-based agents.
+    Automatically handles message conversion to OpenAI format.
+    """
+    def __init__(self, client: Any, model: str = "gpt-4-turbo", system_prompt: str = None):
+        """
+        Args:
+            client: The OpenAI client instance (AsyncOpenAI or OpenAI).
+            model: The model name to use (e.g., "gpt-4-turbo").
+            system_prompt: Optional system instructions for the agent.
+        """
+        self.client = client
+        self.model = model
+        self.system_prompt = system_prompt
+
+    async def call(self, input: AgentInput) -> Union[str, AgentResponse]:
+        # Convert internal message format to OpenAI format
+        # Input messages are already in [{"role": "...", "content": "..."}] format which OpenAI accepts
+        messages = []
+        if self.system_prompt:
+            messages.append({"role": "system", "content": self.system_prompt})
+            
+        messages.extend(input.messages)
+        
+        # Check if client is async or sync
+        if hasattr(self.client, "chat") and hasattr(self.client.chat, "completions"):
+             # Handle AsyncOpenAI vs OpenAI
+            if hasattr(self.client.chat.completions, "create"):
+                 # Modern OpenAI SDK (v1+)
+                if is_async_client(self.client):
+                    completion = await self.client.chat.completions.create(
+                        model=self.model,
+                        messages=messages
+                    )
+                else:
+                    completion = self.client.chat.completions.create(
+                        model=self.model,
+                        messages=messages
+                    )
+                
+                message = completion.choices[0].message
+                content = message.content or ""
+                
+                # Extract tool_calls if present
+                tool_calls = None
+                if hasattr(message, 'tool_calls') and message.tool_calls:
+                    # Convert tool_calls to dict format for serialization
+                    tool_calls = []
+                    for tc in message.tool_calls:
+                        tool_call_dict = {
+                            "id": tc.id,
+                            "type": tc.type,
+                            "function": {
+                                "name": tc.function.name,
+                                "arguments": tc.function.arguments
+                            }
+                        }
+                        tool_calls.append(tool_call_dict)
+                
+                # Return AgentResponse if tool_calls exist, otherwise just content string
+                if tool_calls:
+                    return AgentResponse(content=content, tool_calls=tool_calls)
+                return content
+        
+        raise ValueError("Unsupported OpenAI client. Please provide a valid OpenAI or AsyncOpenAI client.")
+
+def is_async_client(client: Any) -> bool:
+    """Check if the client is an async client."""
+    # Heuristic check for async client
+    return type(client).__name__ == "AsyncOpenAI"
+

fi/simulate/agent/wrappers/websocket.py

@@ -0,0 +1,326 @@
+from __future__ import annotations
+
+import asyncio
+import base64
+import hashlib
+import json
+import os
+import socket
+import time
+from typing import Any, Mapping, Optional
+from urllib.parse import urlparse
+
+from fi.simulate.agent.wrapper import (
+    AgentInput,
+    AgentResponse,
+    SimulationArtifact,
+    SimulationEvent,
+)
+from fi.simulate.agent.wrapper import AgentWrapper
+from fi.simulate.agent.wrappers.http import (
+    _artifact_list,
+    _content_text,
+    _event_list,
+    _optional_mapping,
+    _redacted_endpoint,
+    _tool_call_list,
+    _tool_response_list,
+)
+
+
+class WebSocketAgentWrapper(AgentWrapper):
+    """WebSocket target adapter for local framework transport 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._send_json,
+                request_payload,
+                headers,
+            )
+            response = self._agent_response_from_payload(response_payload)
+        except Exception as exc:
+            error = str(exc)
+            response = AgentResponse(content=f"WebSocket target failed: {exc}")
+
+        latency_ms = round((time.time() - started) * 1000, 4)
+        parsed = urlparse(self.endpoint)
+        trace = {
+            "kind": "external_agent_websocket_trace",
+            "protocol": self.protocol,
+            "endpoint": _redacted_endpoint(self.endpoint),
+            "endpoint_host": parsed.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 status_code == 101,
+            "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_websocket_call",
+                payload=trace,
+            )
+        )
+        response.artifacts.append(
+            SimulationArtifact(
+                type="trace",
+                role="agent",
+                data=trace,
+                metadata={"kind": "external_agent_websocket_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 = list(input.messages)
+        if self.system_prompt:
+            messages = [{"role": "system", "content": self.system_prompt}, *messages]
+        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 = dict(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 _send_json(
+        self,
+        payload: Mapping[str, Any],
+        headers: Mapping[str, str],
+    ) -> tuple[int, dict[str, Any]]:
+        parsed = urlparse(self.endpoint)
+        if parsed.scheme != "ws":
+            raise ValueError("WebSocketAgentWrapper currently supports ws:// endpoints")
+        host = parsed.hostname or ""
+        port = int(parsed.port or 80)
+        path = parsed.path or "/"
+        if parsed.query:
+            path = f"{path}?{parsed.query}"
+        key = base64.b64encode(os.urandom(16)).decode("ascii")
+        with socket.create_connection((host, port), timeout=self.timeout) as sock:
+            sock.settimeout(self.timeout)
+            status = self._handshake(
+                sock,
+                host=host,
+                port=port,
+                path=path,
+                key=key,
+                headers=headers,
+            )
+            _send_text_frame(sock, json.dumps(payload, default=str), mask=True)
+            opcode, text = _read_frame(sock)
+            if opcode == 8:
+                raise ValueError("WebSocket target closed before returning JSON")
+            try:
+                _send_close_frame(sock)
+            except OSError:
+                pass
+        try:
+            response = json.loads(text or "{}")
+        except json.JSONDecodeError as exc:
+            raise ValueError(f"WebSocket target returned non-JSON response: {exc}") from exc
+        if not isinstance(response, dict):
+            raise ValueError("WebSocket target response must be a JSON object")
+        return status, response
+
+    def _handshake(
+        self,
+        sock: socket.socket,
+        *,
+        host: str,
+        port: int,
+        path: str,
+        key: str,
+        headers: Mapping[str, str],
+    ) -> int:
+        request_headers = {
+            "Host": f"{host}:{port}",
+            "Upgrade": "websocket",
+            "Connection": "Upgrade",
+            "Sec-WebSocket-Key": key,
+            "Sec-WebSocket-Version": "13",
+            **dict(headers),
+        }
+        request = "GET " + path + " HTTP/1.1\r\n" + "\r\n".join(
+            f"{name}: {value}" for name, value in request_headers.items()
+        ) + "\r\n\r\n"
+        sock.sendall(request.encode("utf-8"))
+        response = _read_until(sock, b"\r\n\r\n").decode("utf-8", errors="replace")
+        lines = response.split("\r\n")
+        status_line = lines[0] if lines else ""
+        parts = status_line.split()
+        status = int(parts[1]) if len(parts) >= 2 and parts[1].isdigit() else 0
+        response_headers: dict[str, str] = {}
+        for line in lines[1:]:
+            if ":" in line:
+                name, value = line.split(":", 1)
+                response_headers[name.strip().lower()] = value.strip()
+        expected_accept = base64.b64encode(
+            hashlib.sha1(
+                (key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode("ascii")
+            ).digest()
+        ).decode("ascii")
+        if status != 101:
+            raise ValueError(f"WebSocket target returned status {status}")
+        if response_headers.get("sec-websocket-accept") != expected_accept:
+            raise ValueError("WebSocket target returned invalid accept key")
+        return status
+
+    def _agent_response_from_payload(self, payload: Mapping[str, Any]) -> AgentResponse:
+        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 = {
+        "agent_learning_websocket": "fi.alk",
+        "websocket": "fi.alk",
+        "ws": "fi.alk",
+    }
+    protocol = aliases.get(protocol, protocol)
+    if protocol != "fi.alk":
+        raise ValueError("protocol must be fi.alk")
+    return protocol
+
+
+def _read_until(sock: socket.socket, marker: bytes) -> bytes:
+    chunks: list[bytes] = []
+    data = b""
+    while marker not in data:
+        chunk = sock.recv(4096)
+        if not chunk:
+            break
+        chunks.append(chunk)
+        data = b"".join(chunks)
+    return data
+
+
+def _send_text_frame(sock: socket.socket, text: str, *, mask: bool) -> None:
+    payload = text.encode("utf-8")
+    header = bytearray([0x81])
+    length = len(payload)
+    mask_bit = 0x80 if mask else 0
+    if length < 126:
+        header.append(mask_bit | length)
+    elif length <= 0xFFFF:
+        header.extend([mask_bit | 126, *length.to_bytes(2, "big")])
+    else:
+        header.extend([mask_bit | 127, *length.to_bytes(8, "big")])
+    if mask:
+        key = os.urandom(4)
+        masked = bytes(byte ^ key[index % 4] for index, byte in enumerate(payload))
+        sock.sendall(bytes(header) + key + masked)
+    else:
+        sock.sendall(bytes(header) + payload)
+
+
+def _send_close_frame(sock: socket.socket) -> None:
+    sock.sendall(b"\x88\x80" + os.urandom(4))
+
+
+def _read_frame(sock: socket.socket) -> tuple[int, str]:
+    first = _read_exact(sock, 2)
+    opcode = first[0] & 0x0F
+    masked = bool(first[1] & 0x80)
+    length = first[1] & 0x7F
+    if length == 126:
+        length = int.from_bytes(_read_exact(sock, 2), "big")
+    elif length == 127:
+        length = int.from_bytes(_read_exact(sock, 8), "big")
+    mask_key = _read_exact(sock, 4) if masked else b""
+    payload = _read_exact(sock, length) if length else b""
+    if masked:
+        payload = bytes(byte ^ mask_key[index % 4] for index, byte in enumerate(payload))
+    return opcode, payload.decode("utf-8")
+
+
+def _read_exact(sock: socket.socket, size: int) -> bytes:
+    data = bytearray()
+    while len(data) < size:
+        chunk = sock.recv(size - len(data))
+        if not chunk:
+            raise ValueError("WebSocket connection closed unexpectedly")
+        data.extend(chunk)
+    return bytes(data)

fi/simulate/artifacts/__init__.py

@@ -0,0 +1,11 @@
+from .manifest import (
+    ARTIFACT_MANIFEST_SCHEMA_VERSION,
+    ArtifactManifest,
+    ArtifactManifestEntry,
+)
+
+__all__ = [
+    "ARTIFACT_MANIFEST_SCHEMA_VERSION",
+    "ArtifactManifest",
+    "ArtifactManifestEntry",
+]

fi/simulate/artifacts/manifest.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+from pydantic import BaseModel, Field, JsonValue, model_validator
+
+from fi.simulate.agent.wrapper import ArtifactType
+from fi.simulate.evidence import EvidenceClass
+from fi.simulate._hashing import content_hash
+
+ARTIFACT_MANIFEST_SCHEMA_VERSION = "futureagi.artifact-manifest.v1"
+
+
+class ArtifactManifestEntry(BaseModel):
+    artifact_id: str
+    test_case_id: str | None = None
+    type: ArtifactType
+    path: str | None = None
+    uri: str | None = None
+    checksum: str
+    size_bytes: int = Field(ge=0)
+    mime_type: str | None = None
+    codec: str | None = None
+    sample_rate: int | None = Field(default=None, gt=0)
+    channels: int | None = Field(default=None, gt=0)
+    participant_id: str | None = None
+    track_id: str | None = None
+    leg_id: str | None = None
+    start_time_ns: int | None = None
+    end_time_ns: int | None = None
+    evidence_class: EvidenceClass
+    evidence_source_id: str
+    redacted: bool = False
+    metadata: dict[str, JsonValue] = Field(default_factory=dict)
+
+    @model_validator(mode="after")
+    def _validate_location(self) -> "ArtifactManifestEntry":
+        if not self.path and not self.uri:
+            raise ValueError("artifact_location_missing: path or uri is required")
+        if self.path and self.uri:
+            raise ValueError("artifact_location_ambiguous: provide path or uri, not both")
+        if not self.checksum.startswith("sha256:"):
+            raise ValueError("artifact_checksum_invalid: checksum must use sha256:<digest>")
+        return self
+
+
+class ArtifactManifest(BaseModel):
+    schema_version: str = ARTIFACT_MANIFEST_SCHEMA_VERSION
+    run_id: str
+    entries: list[ArtifactManifestEntry] = Field(default_factory=list)
+    manifest_hash: str | None = None
+
+    def content_hash(self) -> str:
+        payload = self.model_dump(exclude={"manifest_hash"}, exclude_none=True)
+        payload["entries"] = sorted(payload["entries"], key=lambda item: item["artifact_id"])
+        return content_hash(payload)
+
+    @model_validator(mode="after")
+    def _stamp_hash(self) -> "ArtifactManifest":
+        expected = self.content_hash()
+        if self.manifest_hash is not None and self.manifest_hash != expected:
+            raise ValueError("artifact_manifest_hash_mismatch")
+        object.__setattr__(self, "manifest_hash", expected)
+        return self