stacklink-sdk 0.2.0__py3-none-any.whl

stacklink/openai_memory.py

@@ -0,0 +1,432 @@
+from __future__ import annotations
+
+import inspect
+import json
+import uuid
+from collections.abc import AsyncIterator, Iterator
+from typing import Any, Callable
+
+from .errors import StacklinkError
+
+
+class StacklinkMemoryLifecycleError(RuntimeError):
+    def __init__(self, message: str, *, phase: str, cause: Exception | None = None):
+        super().__init__(message)
+        self.phase = phase
+        self.__cause__ = cause
+
+
+def with_stacklink_memory(
+    openai_client: Any,
+    *,
+    session: Any,
+    application_key: str = "openai",
+    capture: str = "completed_turns",
+    availability_mode: str = "continue",
+    context_options: dict[str, Any] | None = None,
+    consent: dict[str, Any] | None = None,
+    metadata: dict[str, Any] | None = None,
+    on_diagnostic: Callable[[dict[str, Any]], None] | None = None,
+):
+    """Add automatic Stacklink Memory recall and completed-turn capture to OpenAI.
+
+    The returned object preserves the normal ``responses.create`` and
+    ``chat.completions.create`` call shapes. Both sync and async OpenAI clients
+    are supported. Streaming turns are captured only after the stream reaches a
+    successful terminal event.
+    """
+    lifecycle = _MemoryLifecycle(
+        session=session,
+        application_key=application_key,
+        capture=capture,
+        availability_mode=availability_mode,
+        context_options=context_options or {},
+        consent=consent,
+        metadata=metadata or {},
+        on_diagnostic=on_diagnostic,
+    )
+    responses = _CreateResourceProxy(
+        openai_client.responses,
+        lambda request, original: _responses_create(lifecycle, request, original),
+    )
+    completions = _CreateResourceProxy(
+        openai_client.chat.completions,
+        lambda request, original: _chat_create(lifecycle, request, original),
+    )
+    chat = _AttributeProxy(openai_client.chat, {"completions": completions})
+    return _AttributeProxy(openai_client, {"responses": responses, "chat": chat})
+
+
+class _MemoryLifecycle:
+    def __init__(
+        self,
+        *,
+        session: Any,
+        application_key: str,
+        capture: str,
+        availability_mode: str,
+        context_options: dict[str, Any],
+        consent: dict[str, Any] | None,
+        metadata: dict[str, Any],
+        on_diagnostic: Callable[[dict[str, Any]], None] | None,
+    ):
+        if capture not in {"completed_turns", "off"}:
+            raise ValueError("capture must be 'completed_turns' or 'off'.")
+        if availability_mode not in {"continue", "throw"}:
+            raise ValueError("availability_mode must be 'continue' or 'throw'.")
+        self.session = session
+        self.application_key = application_key
+        self.capture = capture
+        self.availability_mode = availability_mode
+        self.context_options = context_options
+        self.consent = consent
+        self.metadata = metadata
+        self.on_diagnostic = on_diagnostic
+
+    def recall(self, query: str) -> str:
+        if not query.strip():
+            return ""
+        try:
+            response = self.session.memory.context({**self.context_options, "query": query})
+            prompt_text = response.get("promptText") if isinstance(response, dict) else None
+            if not isinstance(prompt_text, str):
+                raise StacklinkMemoryLifecycleError(
+                    "Stacklink Memory context did not include promptText.",
+                    phase="recall",
+                )
+            return prompt_text
+        except Exception as error:
+            self._handle_failure("recall", error)
+            return ""
+
+    def capture_turn(
+        self,
+        *,
+        user_text: str,
+        assistant_text: str,
+        framework: str,
+        response_id: str | None,
+    ) -> None:
+        if self.capture == "off" or not user_text.strip() or not assistant_text.strip():
+            return
+        content = json.dumps(
+            {
+                "schemaVersion": "stacklink.conversation.turn.v1",
+                "user": {"role": "user", "content": user_text},
+                "assistant": {"role": "assistant", "content": assistant_text},
+            },
+            separators=(",", ":"),
+            ensure_ascii=False,
+        )
+        body: dict[str, Any] = {
+            "applicationKey": self.application_key,
+            "turnId": response_id or str(uuid.uuid4()),
+            "content": content,
+            "metadata": {
+                **self.metadata,
+                "framework": framework,
+                "responseId": response_id,
+                "captureMode": "completed_turns",
+            },
+        }
+        if self.consent is not None:
+            body["consent"] = self.consent
+        try:
+            self.session.memory.record_conversation_turn(body)
+        except Exception as error:
+            self._handle_failure("capture", error)
+
+    def _handle_failure(self, phase: str, error: Exception) -> None:
+        if self.availability_mode == "continue" and _is_availability_failure(error):
+            if self.on_diagnostic is not None:
+                self.on_diagnostic(
+                    {
+                        "phase": phase,
+                        "code": "memory_unavailable",
+                        "message": (
+                            "The model call continued without Stacklink Memory context."
+                            if phase == "recall"
+                            else "The model response completed, but Stacklink Memory capture is unavailable."
+                        ),
+                        "error": error,
+                    }
+                )
+            return
+        if isinstance(error, StacklinkMemoryLifecycleError):
+            raise error
+        raise StacklinkMemoryLifecycleError(
+            (
+                "Stacklink Memory context could not be assembled."
+                if phase == "recall"
+                else "The completed turn could not be accepted by Stacklink Memory."
+            ),
+            phase=phase,
+            cause=error,
+        ) from error
+
+
+def _responses_create(lifecycle: _MemoryLifecycle, request: dict[str, Any], original: Callable[..., Any]):
+    user_text = _responses_user_text(request.get("input"))
+    prompt_text = lifecycle.recall(user_text)
+    next_request = dict(request)
+    if prompt_text:
+        memory_instructions = f"Stacklink Memory context\n\n{prompt_text}"
+        existing = next_request.get("instructions")
+        next_request["instructions"] = (
+            f"{existing}\n\n{memory_instructions}"
+            if isinstance(existing, str) and existing
+            else memory_instructions
+        )
+    response = original(**next_request)
+
+    def transform(value: Any):
+        if request.get("stream") is True:
+            observer = _ResponsesStreamObserver(lifecycle, user_text)
+            return _wrap_stream(value, observer)
+        if _value(value, "status") == "completed":
+            lifecycle.capture_turn(
+                user_text=user_text,
+                assistant_text=_string(_value(value, "output_text")),
+                framework="openai-python-responses",
+                response_id=_optional_string(_value(value, "id")),
+            )
+        return value
+
+    return _transform_possible_awaitable(response, transform)
+
+
+def _chat_create(lifecycle: _MemoryLifecycle, request: dict[str, Any], original: Callable[..., Any]):
+    user_text = _chat_user_text(request.get("messages"))
+    prompt_text = lifecycle.recall(user_text)
+    next_request = dict(request)
+    if prompt_text:
+        messages = list(next_request.get("messages") or [])
+        insertion = len(messages)
+        for index, message in enumerate(messages):
+            role = _value(message, "role")
+            if role not in {"system", "developer"}:
+                insertion = index
+                break
+        messages.insert(insertion, {"role": "system", "content": f"Stacklink Memory context\n\n{prompt_text}"})
+        next_request["messages"] = messages
+    response = original(**next_request)
+
+    def transform(value: Any):
+        if request.get("stream") is True:
+            observer = _ChatStreamObserver(lifecycle, user_text)
+            return _wrap_stream(value, observer)
+        choices = _value(value, "choices")
+        choice = choices[0] if isinstance(choices, list) and choices else None
+        finish_reason = _value(choice, "finish_reason")
+        if finish_reason in {"stop", "length"}:
+            lifecycle.capture_turn(
+                user_text=user_text,
+                assistant_text=_string(_value(_value(choice, "message"), "content")),
+                framework="openai-python-chat-completions",
+                response_id=_optional_string(_value(value, "id")),
+            )
+        return value
+
+    return _transform_possible_awaitable(response, transform)
+
+
+async def _await_and_transform(awaitable: Any, transform: Callable[[Any], Any]):
+    return transform(await awaitable)
+
+
+def _transform_possible_awaitable(value: Any, transform: Callable[[Any], Any]):
+    if inspect.isawaitable(value):
+        return _await_and_transform(value, transform)
+    return transform(value)
+
+
+def _wrap_stream(source: Any, observer: "_StreamObserver"):
+    if hasattr(source, "__aiter__"):
+        return _ObservedAsyncIterator(source, observer)
+    if hasattr(source, "__iter__"):
+        return _ObservedIterator(source, observer)
+    raise TypeError("OpenAI returned a stream without a sync or async iterator contract.")
+
+
+class _StreamObserver:
+    def item(self, value: Any) -> None:
+        raise NotImplementedError
+
+    def complete(self) -> None:
+        raise NotImplementedError
+
+
+class _ResponsesStreamObserver(_StreamObserver):
+    def __init__(self, lifecycle: _MemoryLifecycle, user_text: str):
+        self.lifecycle = lifecycle
+        self.user_text = user_text
+        self.assistant_text = ""
+        self.response_id: str | None = None
+        self.completed = False
+
+    def item(self, event: Any) -> None:
+        event_type = _value(event, "type")
+        if event_type == "response.output_text.delta":
+            self.assistant_text += _string(_value(event, "delta"))
+        elif event_type == "response.completed":
+            self.completed = True
+            response = _value(event, "response")
+            self.response_id = _optional_string(_value(response, "id"))
+            if not self.assistant_text:
+                self.assistant_text = _string(_value(response, "output_text"))
+
+    def complete(self) -> None:
+        if self.completed:
+            self.lifecycle.capture_turn(
+                user_text=self.user_text,
+                assistant_text=self.assistant_text,
+                framework="openai-python-responses",
+                response_id=self.response_id,
+            )
+
+
+class _ChatStreamObserver(_StreamObserver):
+    def __init__(self, lifecycle: _MemoryLifecycle, user_text: str):
+        self.lifecycle = lifecycle
+        self.user_text = user_text
+        self.assistant_text = ""
+        self.response_id: str | None = None
+        self.finish_reason: str | None = None
+
+    def item(self, chunk: Any) -> None:
+        self.response_id = self.response_id or _optional_string(_value(chunk, "id"))
+        choices = _value(chunk, "choices")
+        choice = choices[0] if isinstance(choices, list) and choices else None
+        self.assistant_text += _string(_value(_value(choice, "delta"), "content"))
+        self.finish_reason = _optional_string(_value(choice, "finish_reason")) or self.finish_reason
+
+    def complete(self) -> None:
+        if self.finish_reason in {"stop", "length"}:
+            self.lifecycle.capture_turn(
+                user_text=self.user_text,
+                assistant_text=self.assistant_text,
+                framework="openai-python-chat-completions",
+                response_id=self.response_id,
+            )
+
+
+class _ObservedIterator(Iterator[Any]):
+    def __init__(self, source: Any, observer: _StreamObserver):
+        self.source = source
+        self.iterator = iter(source)
+        self.observer = observer
+        self.finished = False
+
+    def __next__(self):
+        try:
+            item = next(self.iterator)
+        except StopIteration:
+            if not self.finished:
+                self.finished = True
+                self.observer.complete()
+            raise
+        self.observer.item(item)
+        return item
+
+    def __getattr__(self, name: str):
+        return getattr(self.source, name)
+
+
+class _ObservedAsyncIterator(AsyncIterator[Any]):
+    def __init__(self, source: Any, observer: _StreamObserver):
+        self.source = source
+        self.iterator = source.__aiter__()
+        self.observer = observer
+        self.finished = False
+
+    async def __anext__(self):
+        try:
+            item = await self.iterator.__anext__()
+        except StopAsyncIteration:
+            if not self.finished:
+                self.finished = True
+                self.observer.complete()
+            raise
+        self.observer.item(item)
+        return item
+
+    def __getattr__(self, name: str):
+        return getattr(self.source, name)
+
+
+class _AttributeProxy:
+    def __init__(self, target: Any, overrides: dict[str, Any]):
+        self._target = target
+        self._overrides = overrides
+
+    def __getattr__(self, name: str):
+        if name in self._overrides:
+            return self._overrides[name]
+        return getattr(self._target, name)
+
+
+class _CreateResourceProxy(_AttributeProxy):
+    def __init__(self, target: Any, create: Callable[[dict[str, Any], Callable[..., Any]], Any]):
+        super().__init__(target, {})
+        self._create = create
+
+    def create(self, **request: Any):
+        return self._create(request, self._target.create)
+
+
+def _responses_user_text(value: Any) -> str:
+    if isinstance(value, str):
+        return value
+    if not isinstance(value, list):
+        return ""
+    for item in reversed(value):
+        if _value(item, "role") != "user":
+            continue
+        content = _value(item, "content")
+        if isinstance(content, str):
+            return content
+        if isinstance(content, list):
+            return "\n".join(
+                _string(_value(part, "text"))
+                for part in content
+                if _value(part, "type") == "input_text"
+            )
+    return ""
+
+
+def _chat_user_text(value: Any) -> str:
+    if not isinstance(value, list):
+        return ""
+    for message in reversed(value):
+        if _value(message, "role") != "user":
+            continue
+        content = _value(message, "content")
+        if isinstance(content, str):
+            return content
+        if isinstance(content, list):
+            return "\n".join(
+                _string(_value(part, "text"))
+                for part in content
+                if _value(part, "type") == "text"
+            )
+    return ""
+
+
+def _is_availability_failure(error: Exception) -> bool:
+    if isinstance(error, StacklinkError) and error.status_code is not None:
+        return error.status_code == 408 or error.status_code >= 500
+    return isinstance(error, (ConnectionError, TimeoutError, OSError))
+
+
+def _value(value: Any, name: str):
+    if isinstance(value, dict):
+        return value.get(name)
+    return getattr(value, name, None)
+
+
+def _string(value: Any) -> str:
+    return value if isinstance(value, str) else ""
+
+
+def _optional_string(value: Any) -> str | None:
+    return value if isinstance(value, str) and value else None

stacklink/projects.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+class ProjectsResource:
+    def __init__(self, client):
+        self.client = client
+
+    def list(self, **query):
+        return self.client.request("GET", "/v1/projects", query=query)
+
+    def get(self, project_id: str):
+        return self.client.request("GET", f"/v1/projects/{project_id}")
+
+    def create(self, body: dict[str, Any]):
+        return self.client.request("POST", "/v1/projects", json_body=body, registration=True)
+
+    def reset_auth_configs(self, project_id: str):
+        return self.client.request("POST", f"/v1/projects/{project_id}/auth-configs/reset", registration=True)

stacklink/sessions.py

@@ -0,0 +1,223 @@
+from __future__ import annotations
+
+from typing import Any, Mapping, TypedDict
+
+from .memory import SessionMemoryResource
+from .types import ProjectProvider
+
+
+class SessionToolDefinition(TypedDict):
+    name: str
+    description: str
+    inputSchema: dict[str, Any]
+
+
+class SessionTool:
+    def __init__(self, session: "Session", definition: SessionToolDefinition):
+        self.session = session
+        self.name = definition["name"]
+        self.description = definition.get("description", "")
+        self.parameters = definition.get("inputSchema", {})
+
+    def call(self, arguments: Mapping[str, Any] | None = None):
+        return self.session.client.mcp.call_session_tool_result(self.session.id, self.name, dict(arguments or {}))
+
+    def __call__(self, **arguments: Any):
+        return self.call(arguments)
+
+
+class Session:
+    def __init__(
+        self,
+        client,
+        session_id: str,
+        *,
+        session_data: dict[str, Any] | None = None,
+        tool_projection: dict[str, Any] | None = None,
+    ):
+        self.client = client
+        self.id = session_id
+        self.data = dict(session_data or {})
+        self._tool_projection = dict(tool_projection or {})
+        self.memory = SessionMemoryResource(client.memory, session_id)
+        self.sandbox = client.workbench.bind(session_id)
+
+    def get(self):
+        return self.client.sessions.get(self.id)
+
+    def update(self, body: dict[str, Any]):
+        response = self.client.sessions.update(self.id, body)
+        self.data = dict(response.get("session", self.data)) if isinstance(response, dict) else self.data
+        self.refresh_tools()
+        return response
+
+    def set_provider_mutation_mode(
+        self,
+        mode: str,
+        *,
+        expected_config_version: int,
+    ):
+        if mode not in {"read_only", "ask", "full_access"}:
+            raise ValueError("mode must be read_only, ask, or full_access")
+        return self.update({
+            "expectedConfigVersion": expected_config_version,
+            "providerMutationMode": mode,
+        })
+
+    def close(self):
+        return self.client.sessions.close(self.id)
+
+    def fail(self, body: dict[str, Any] | None = None):
+        return self.client.sessions.fail(self.id, body)
+
+    def trace(self, **query):
+        return self.client.sessions.trace(self.id, **query)
+
+    def config_history(self, **query):
+        return self.client.sessions.config_history(self.id, **query)
+
+    def refresh_tools(self) -> dict[str, Any]:
+        self._tool_projection = self.client.sessions.meta_tools(self.id)
+        return self._tool_projection
+
+    @property
+    def tool_contract_version(self) -> str:
+        self._ensure_tools()
+        return str(self._tool_projection.get("contractVersion", "unknown"))
+
+    def tool_registry(self) -> list[SessionToolDefinition]:
+        self._ensure_tools()
+        tools = self._tool_projection.get("tools", [])
+        return [dict(tool) for tool in tools if isinstance(tool, dict) and isinstance(tool.get("name"), str)]
+
+    def tools(self) -> dict[str, SessionTool]:
+        return {definition["name"]: SessionTool(self, definition) for definition in self.tool_registry()}
+
+    def mcp_connection(self) -> dict[str, Any]:
+        descriptor = self.data.get("mcp")
+        if not isinstance(descriptor, dict) or not isinstance(descriptor.get("url"), str):
+            response = self.get()
+            session_data = response.get("session", {}) if isinstance(response, dict) else {}
+            self.data = dict(session_data)
+            descriptor = self.data.get("mcp")
+        if not isinstance(descriptor, dict) or not isinstance(descriptor.get("url"), str):
+            raise ValueError("This Stacklink session does not include an MCP connection descriptor.")
+        headers = descriptor.get("headers")
+        return {
+            "url": descriptor["url"],
+            "headers": {
+                **(headers if isinstance(headers, dict) else {}),
+                "Authorization": f"Bearer {self.client.api_key}",
+            },
+        }
+
+    def langchain_tools(self) -> list[dict[str, Any]]:
+        return [
+            {
+                "name": tool.name,
+                "description": tool.description,
+                "schema": tool.parameters,
+                "invoke": tool.call,
+                "func": tool.call,
+            }
+            for tool in self.tools().values()
+        ]
+
+    def _ensure_tools(self) -> None:
+        if not isinstance(self._tool_projection.get("tools"), list):
+            self.refresh_tools()
+
+
+class SessionsResource:
+    def __init__(self, client):
+        self.client = client
+
+    def create(self, body: dict[str, Any]):
+        return self.client.request("POST", "/v1/sessions", json_body=body)
+
+    def create_bound(self, body: dict[str, Any]) -> Session:
+        response = self.create(body)
+        session_data = response.get("session", {}) if isinstance(response, dict) else {}
+        session_id = str(session_data["id"])
+        return Session(
+            self.client,
+            session_id,
+            session_data=session_data,
+            tool_projection=self.meta_tools(session_id),
+        )
+
+    def bind(self, session_id: str) -> Session:
+        return Session(self.client, session_id)
+
+    def use(self, session_id: str) -> Session:
+        response = self.get(session_id)
+        session_data = response.get("session", {}) if isinstance(response, dict) else {}
+        return Session(
+            self.client,
+            session_id,
+            session_data=session_data,
+            tool_projection=self.meta_tools(session_id),
+        )
+
+    def list(self, **query):
+        return self.client.request("GET", "/v1/sessions", query=query)
+
+    def get(self, session_id: str):
+        return self.client.request("GET", f"/v1/sessions/{session_id}")
+
+    def update(self, session_id: str, body: dict[str, Any]):
+        response = self.client.request("PATCH", f"/v1/sessions/{session_id}", json_body=body)
+        self.client.clear_cache(f"session:{session_id}:")
+        return response
+
+    def meta_tools(self, session_id: str):
+        return self.client.request("GET", f"/v1/sessions/{session_id}/meta-tools")
+
+    def close(self, session_id: str):
+        return self.client.request("POST", f"/v1/sessions/{session_id}/close")
+
+    def fail(self, session_id: str, body: dict[str, Any] | None = None):
+        return self.client.request("POST", f"/v1/sessions/{session_id}/fail", json_body=body or {})
+
+    def router_facts(self, session_id: str, **query):
+        return self.client.request("GET", f"/v1/sessions/{session_id}/router-facts", query=query)
+
+    def workbench(self, session_id: str):
+        return self.client.request("GET", f"/v1/sessions/{session_id}/workbench")
+
+    def trace(self, session_id: str, **query):
+        return self.client.request("GET", f"/v1/sessions/{session_id}/trace", query=query)
+
+    def config_history(self, session_id: str, **query):
+        return self.client.request("GET", f"/v1/sessions/{session_id}/config-history", query=query)
+
+    def toolkits(self, session_id: str) -> dict[str, list[ProjectProvider]]:
+        return self.client.cached_request(
+            f"session:{session_id}:toolkits",
+            "GET",
+            f"/v1/sessions/{session_id}/toolkits",
+            ttl_seconds=300,
+        )
+
+    def search_tools(self, session_id: str, body: dict[str, Any], *, cache: str = "default"):
+        return self.client.tools.search_for_session(session_id, body, cache=cache)
+
+    def get_tool_schemas(self, session_id: str, body: dict[str, Any], *, cache: str = "default"):
+        return self.client.tools.schemas_for_session(session_id, body, cache=cache)
+
+    def prepare_tools(self, session_id: str, **kwargs):
+        return self.client.tools.prepare_for_session(session_id, **kwargs)
+
+    def prepare_runtime_tools(self, session_id: str, **kwargs):
+        return self.client.tools.prepare_runtime_for_session(session_id, **kwargs)
+
+    def stream_events(self, session_id: str, **query):
+        return self.client.request_sse("GET", f"/v1/sessions/{session_id}/events/stream", query=query)
+
+    def authorize_connection(self, session_id: str, body: dict[str, Any]):
+        return self.client.request("POST", f"/v1/sessions/{session_id}/connections/authorize", json_body=body)
+
+    def mcp_url(self, session: dict[str, Any] | Session) -> str | None:
+        value = session.data if isinstance(session, Session) else session
+        mcp = value.get("mcp")
+        return mcp.get("url") if isinstance(mcp, dict) else None