init
This commit is contained in:
@@ -0,0 +1,153 @@
|
||||
/**
|
||||
* AtlasClient — stateless HTTP/SSE transport layer.
|
||||
*
|
||||
* Single responsibility: send chat payloads to the Cognite AI agent API
|
||||
* and parse the response (JSON or Server-Sent Events).
|
||||
*/
|
||||
|
||||
import type { CogniteClient } from '@cognite/sdk';
|
||||
import type { Agent, ChatPayload, RawAgentResponse, StreamCallbacks } from './types';
|
||||
|
||||
const CDF_API_VERSION = 'alpha';
|
||||
const AGENTS_API_VERSION = 'beta';
|
||||
|
||||
export class AtlasClient {
|
||||
private readonly client: CogniteClient;
|
||||
|
||||
constructor(client: CogniteClient) {
|
||||
this.client = client;
|
||||
}
|
||||
|
||||
async listAgents(): Promise<Agent[]> {
|
||||
const { data } = await this.client.get<{ items: Agent[] }>(
|
||||
`/api/v1/projects/${this.client.project}/ai/agents`,
|
||||
{ headers: { 'cdf-version': AGENTS_API_VERSION } },
|
||||
);
|
||||
return data.items;
|
||||
}
|
||||
|
||||
async getAgentById(externalId: string): Promise<Agent | null> {
|
||||
const { data } = await this.client.post<{ items: Agent[] }>(
|
||||
`/api/v1/projects/${this.client.project}/ai/agents/byids`,
|
||||
{ data: { items: [{ externalId }] }, headers: { 'cdf-version': AGENTS_API_VERSION } },
|
||||
);
|
||||
return data.items[0] ?? null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Post a chat payload and parse the response (JSON or SSE).
|
||||
* @param agentExternalId — used as a fallback identifier when the SSE result event omits agent IDs.
|
||||
*/
|
||||
async post(
|
||||
payload: ChatPayload,
|
||||
agentExternalId: string,
|
||||
callbacks?: StreamCallbacks,
|
||||
signal?: AbortSignal,
|
||||
): Promise<RawAgentResponse> {
|
||||
const url = `${this.client.getBaseUrl()}/api/v1/projects/${this.client.project}/ai/internal/agents/chat`;
|
||||
|
||||
const response = await fetch(url, {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
...this.client.getDefaultRequestHeaders(),
|
||||
'Content-Type': 'application/json',
|
||||
'cdf-version': CDF_API_VERSION,
|
||||
},
|
||||
body: JSON.stringify(payload),
|
||||
signal,
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const errorData = await response.json().catch(() => ({}));
|
||||
throw new Error(
|
||||
`Agent chat API error: ${response.status} - ${JSON.stringify(errorData)}`,
|
||||
);
|
||||
}
|
||||
|
||||
const contentType = response.headers.get('content-type') || '';
|
||||
|
||||
if (
|
||||
contentType.includes('text/event-stream') ||
|
||||
contentType.includes('text/plain')
|
||||
) {
|
||||
return this.parseSSE(response, agentExternalId, callbacks);
|
||||
}
|
||||
|
||||
return await response.json();
|
||||
}
|
||||
|
||||
/**
|
||||
* Parse a Server-Sent Events streaming response.
|
||||
*/
|
||||
private async parseSSE(
|
||||
response: Response,
|
||||
agentExternalId: string,
|
||||
callbacks?: StreamCallbacks,
|
||||
): Promise<RawAgentResponse> {
|
||||
const reader = response.body?.getReader();
|
||||
if (!reader) {
|
||||
throw new Error('Response body is not readable');
|
||||
}
|
||||
|
||||
const decoder = new TextDecoder();
|
||||
let buffer = '';
|
||||
let finalResponse: RawAgentResponse | null = null;
|
||||
|
||||
const processLine = (line: string) => {
|
||||
if (!line.startsWith('data: ')) return;
|
||||
|
||||
const dataStr = line.slice(6).trim();
|
||||
if (dataStr === '[DONE]') return '[DONE]' as const;
|
||||
|
||||
try {
|
||||
const data = JSON.parse(dataStr);
|
||||
const sseResponse = data.response;
|
||||
|
||||
if (!sseResponse) return;
|
||||
|
||||
if (sseResponse.type === 'progress' && sseResponse.content) {
|
||||
callbacks?.onProgress?.(sseResponse.content);
|
||||
} else if (
|
||||
sseResponse.type === 'responseChunk' &&
|
||||
sseResponse.content
|
||||
) {
|
||||
callbacks?.onChunk?.(sseResponse.content);
|
||||
} else if (sseResponse.type === 'result') {
|
||||
finalResponse = {
|
||||
agentId: data.agentId || agentExternalId,
|
||||
agentExternalId: data.agentExternalId || agentExternalId,
|
||||
response: {
|
||||
type: 'result',
|
||||
cursor: sseResponse.cursor,
|
||||
messages: sseResponse.messages || [],
|
||||
},
|
||||
};
|
||||
}
|
||||
} catch {
|
||||
// Skip unparseable SSE lines
|
||||
}
|
||||
};
|
||||
|
||||
try {
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
|
||||
if (done) break;
|
||||
|
||||
buffer += decoder.decode(value, { stream: true });
|
||||
const lines = buffer.split('\n');
|
||||
buffer = lines.pop() || '';
|
||||
|
||||
if (lines.some((line) => processLine(line) === '[DONE]')) break;
|
||||
}
|
||||
} finally {
|
||||
reader.releaseLock();
|
||||
}
|
||||
|
||||
if (!finalResponse) {
|
||||
throw new Error('No result response received from streaming API');
|
||||
}
|
||||
|
||||
return finalResponse;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
/**
|
||||
* Atlas Agent Client — public API.
|
||||
*
|
||||
* Core library for communicating with Cognite Atlas AI agents.
|
||||
* Self-contained, zero imports from outside this directory (except external packages).
|
||||
*
|
||||
* React hook is a separate import path:
|
||||
* import { useAtlasChat } from '@cognite/dune-utils/atlas-agent/react';
|
||||
*/
|
||||
|
||||
// Core
|
||||
export { AtlasSession } from './session';
|
||||
export { AtlasClient } from './client';
|
||||
|
||||
// TypeBox re-exports for convenience
|
||||
export { Type } from '@sinclair/typebox';
|
||||
export type { Static, TSchema } from '@sinclair/typebox';
|
||||
|
||||
// Types
|
||||
export type {
|
||||
Agent,
|
||||
AgentToolConfig,
|
||||
AtlasTool,
|
||||
AtlasToolResult,
|
||||
AtlasResponse,
|
||||
AtlasSessionConfig,
|
||||
ToolCall,
|
||||
StreamCallbacks,
|
||||
ApiToolDefinition,
|
||||
PythonRuntime,
|
||||
} from './types';
|
||||
@@ -0,0 +1,146 @@
|
||||
/**
|
||||
* usePyodideRuntime — React hook for managing PyodideRuntime lifecycle.
|
||||
*
|
||||
* Separate entry point so the core atlas-agent bundle stays Pyodide-free.
|
||||
*/
|
||||
|
||||
import { useState, useRef, useEffect } from 'react';
|
||||
import type { CogniteClient } from '@cognite/sdk';
|
||||
import { getGlobalPyodideRuntime } from './pyodide-runtime';
|
||||
import type { PyodideRuntimeConfig } from './pyodide-runtime';
|
||||
import type { PythonRuntime } from './types';
|
||||
|
||||
// ============================================================================
|
||||
// Types
|
||||
// ============================================================================
|
||||
|
||||
export interface PyodideProgress {
|
||||
stage: string;
|
||||
percent: number;
|
||||
}
|
||||
|
||||
export interface UsePyodideRuntimeOptions {
|
||||
/** The `loadPyodide` function from the `pyodide` package. */
|
||||
loadPyodide: PyodideRuntimeConfig['loadPyodide'];
|
||||
/** CogniteClient for SDK credential injection. `null` disables initialization. */
|
||||
client: CogniteClient | null;
|
||||
/** Additional Python packages to install via micropip. */
|
||||
requirements?: string[];
|
||||
/** CDN URL for Pyodide files. */
|
||||
cdnUrl?: string;
|
||||
}
|
||||
|
||||
export interface UsePyodideRuntimeReturn {
|
||||
/** The initialized runtime, or undefined if not yet ready. */
|
||||
runtime: PythonRuntime | undefined;
|
||||
/** True while Pyodide is loading / initializing. */
|
||||
loading: boolean;
|
||||
/** Error message if initialization failed. */
|
||||
error: string | null;
|
||||
/** Current initialization progress. */
|
||||
progress: PyodideProgress;
|
||||
/** Convenience: true when runtime is ready to use. */
|
||||
isReady: boolean;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Hook
|
||||
// ============================================================================
|
||||
|
||||
const DEFAULT_BASE_URL = 'https://api.cognitedata.com';
|
||||
|
||||
/**
|
||||
* Manages PyodideRuntime initialization lifecycle.
|
||||
*
|
||||
* Loads Pyodide, installs packages, injects Cognite SDK credentials,
|
||||
* and returns a ready-to-use `PythonRuntime` with loading/error state.
|
||||
*
|
||||
* ```tsx
|
||||
* import { loadPyodide } from 'pyodide';
|
||||
*
|
||||
* const { runtime, loading, progress, isReady } = usePyodideRuntime({
|
||||
* loadPyodide,
|
||||
* client: sdk,
|
||||
* requirements: ['pandas', 'numpy'],
|
||||
* });
|
||||
* ```
|
||||
*/
|
||||
export function usePyodideRuntime(
|
||||
options: UsePyodideRuntimeOptions,
|
||||
): UsePyodideRuntimeReturn {
|
||||
const { client } = options;
|
||||
|
||||
const [runtime, setRuntime] = useState<PythonRuntime>();
|
||||
const [loading, setLoading] = useState(true);
|
||||
const [error, setError] = useState<string | null>(null);
|
||||
const [progress, setProgress] = useState<PyodideProgress>({ stage: '', percent: 0 });
|
||||
|
||||
// Refs for values that shouldn't trigger re-initialization
|
||||
const loadPyodideRef = useRef(options.loadPyodide);
|
||||
const requirementsRef = useRef(options.requirements);
|
||||
const cdnUrlRef = useRef(options.cdnUrl);
|
||||
loadPyodideRef.current = options.loadPyodide;
|
||||
requirementsRef.current = options.requirements;
|
||||
cdnUrlRef.current = options.cdnUrl;
|
||||
|
||||
useEffect(() => {
|
||||
if (!client) {
|
||||
setLoading(false);
|
||||
return;
|
||||
}
|
||||
|
||||
let mounted = true;
|
||||
|
||||
(async () => {
|
||||
try {
|
||||
setLoading(true);
|
||||
setError(null);
|
||||
|
||||
const instance = getGlobalPyodideRuntime({
|
||||
loadPyodide: loadPyodideRef.current,
|
||||
requirements: requirementsRef.current,
|
||||
cdnUrl: cdnUrlRef.current,
|
||||
onProgress: (stage, percent) => {
|
||||
if (mounted) setProgress({ stage, percent });
|
||||
},
|
||||
});
|
||||
|
||||
if (!instance.isInitialized) {
|
||||
if (mounted) setProgress({ stage: 'Initializing...', percent: 0 });
|
||||
|
||||
const headers = client.getDefaultRequestHeaders();
|
||||
const token = headers.Authorization?.split(' ')[1] ?? '';
|
||||
|
||||
await instance.initialize({
|
||||
project: client.project,
|
||||
baseUrl: client.getBaseUrl?.() ?? DEFAULT_BASE_URL,
|
||||
token,
|
||||
});
|
||||
}
|
||||
|
||||
if (mounted) {
|
||||
setRuntime(instance);
|
||||
setProgress({ stage: 'Ready', percent: 100 });
|
||||
}
|
||||
} catch (err) {
|
||||
if (mounted) {
|
||||
setError(err instanceof Error ? err.message : String(err));
|
||||
}
|
||||
} finally {
|
||||
if (mounted) setLoading(false);
|
||||
}
|
||||
})();
|
||||
|
||||
return () => {
|
||||
mounted = false;
|
||||
};
|
||||
}, [client]);
|
||||
|
||||
return {
|
||||
runtime,
|
||||
loading,
|
||||
error,
|
||||
progress,
|
||||
isReady: !loading && !error && runtime !== undefined,
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,295 @@
|
||||
/**
|
||||
* PyodideRuntime — browser-based Python execution via Pyodide.
|
||||
*
|
||||
* Wraps Pyodide loading, package installation, Cognite SDK setup,
|
||||
* and localStorage caching into a clean PythonRuntime implementation.
|
||||
*
|
||||
* The consumer owns the 'pyodide' npm package — they pass `loadPyodide`
|
||||
* as a config parameter so this module has no hard dependency on it.
|
||||
*/
|
||||
|
||||
import type { PythonRuntime } from './types';
|
||||
|
||||
// ============================================================================
|
||||
// Minimal Pyodide Interfaces (avoids 'pyodide' package dependency)
|
||||
// ============================================================================
|
||||
|
||||
interface PyodideGlobals {
|
||||
get(name: string): unknown;
|
||||
set(name: string, value: unknown): void;
|
||||
}
|
||||
|
||||
/** Subset of PyodideInterface that this module uses. */
|
||||
export interface PyodideInstance {
|
||||
loadPackage(packages: string[]): Promise<void>;
|
||||
runPython(code: string): unknown;
|
||||
runPythonAsync(code: string): Promise<unknown>;
|
||||
globals: PyodideGlobals;
|
||||
pyimport(name: string): unknown;
|
||||
}
|
||||
|
||||
interface Micropip {
|
||||
install(packages: string | string[]): Promise<void>;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Constants
|
||||
// ============================================================================
|
||||
|
||||
const DEFAULT_CDN_URL = 'https://cdn.jsdelivr.net/pyodide/v0.29.3/full/';
|
||||
const CACHE_KEY = 'dune_pyodide_initialized';
|
||||
const CACHE_VERSION = 'v1';
|
||||
|
||||
// ============================================================================
|
||||
// Config Types
|
||||
// ============================================================================
|
||||
|
||||
export interface PyodideRuntimeConfig {
|
||||
/** The `loadPyodide` function from the `pyodide` package. */
|
||||
loadPyodide: (options: { indexURL: string }) => Promise<PyodideInstance>;
|
||||
/** CDN URL for Pyodide files. Defaults to jsDelivr v0.29.3. */
|
||||
cdnUrl?: string;
|
||||
/** Additional Python packages to install via micropip. */
|
||||
requirements?: string[];
|
||||
/** Progress callback for initialization stages. */
|
||||
onProgress?: (stage: string, percent: number) => void;
|
||||
}
|
||||
|
||||
export interface PyodideSDKConfig {
|
||||
project: string;
|
||||
baseUrl: string;
|
||||
token: string;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Python Utility Code (injected at init)
|
||||
// ============================================================================
|
||||
|
||||
const PYTHON_UTILS = `
|
||||
import json
|
||||
|
||||
def _serialize_cognite_object(obj, depth=0):
|
||||
if depth > 10: return str(obj)
|
||||
for attr in ('dump', 'as_dict'):
|
||||
fn = getattr(obj, attr, None)
|
||||
if fn:
|
||||
try: return fn()
|
||||
except: pass
|
||||
if isinstance(obj, dict):
|
||||
return {k: _serialize_cognite_object(v, depth+1) for k, v in obj.items()}
|
||||
if isinstance(obj, (list, tuple)):
|
||||
return [_serialize_cognite_object(i, depth+1) for i in obj]
|
||||
if isinstance(obj, (str, int, float, bool, type(None))):
|
||||
return obj
|
||||
d = getattr(obj, '__dict__', None)
|
||||
if d is not None:
|
||||
try: return _serialize_cognite_object(d, depth+1)
|
||||
except: pass
|
||||
return str(obj)
|
||||
|
||||
def as_json_string(value):
|
||||
return json.dumps(_serialize_cognite_object(value))
|
||||
`;
|
||||
|
||||
// ============================================================================
|
||||
// Cache Helpers
|
||||
// ============================================================================
|
||||
|
||||
function isCacheValid(): boolean {
|
||||
try {
|
||||
return localStorage.getItem(CACHE_KEY) === CACHE_VERSION;
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
function markCacheValid(): void {
|
||||
try {
|
||||
localStorage.setItem(CACHE_KEY, CACHE_VERSION);
|
||||
} catch {
|
||||
/* localStorage unavailable */
|
||||
}
|
||||
}
|
||||
|
||||
/** Clear the Pyodide package cache — forces re-download on next init. */
|
||||
export function clearPyodideCache(): void {
|
||||
try {
|
||||
localStorage.removeItem(CACHE_KEY);
|
||||
} catch {
|
||||
/* localStorage unavailable */
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// PyProxy Detection (structural — avoids importing from 'pyodide')
|
||||
// ============================================================================
|
||||
|
||||
function isPyProxy(value: unknown): value is { destroy(): void } {
|
||||
return (
|
||||
value != null &&
|
||||
typeof value === 'object' &&
|
||||
'destroy' in value &&
|
||||
typeof (value as Record<string, unknown>).destroy === 'function'
|
||||
);
|
||||
}
|
||||
|
||||
function destroyIfPyProxy(value: unknown): void {
|
||||
if (isPyProxy(value)) {
|
||||
value.destroy();
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// PyodideRuntime Class
|
||||
// ============================================================================
|
||||
|
||||
/**
|
||||
* PythonRuntime backed by Pyodide. Handles loading, package installation,
|
||||
* Cognite SDK credential injection, caching, and PyProxy conversion.
|
||||
*/
|
||||
export class PyodideRuntime implements PythonRuntime {
|
||||
private pyodide?: PyodideInstance;
|
||||
private micropip?: Micropip;
|
||||
private _initialized = false;
|
||||
private readonly config: PyodideRuntimeConfig;
|
||||
|
||||
constructor(config: PyodideRuntimeConfig) {
|
||||
this.config = config;
|
||||
}
|
||||
|
||||
get isInitialized(): boolean {
|
||||
return this._initialized;
|
||||
}
|
||||
|
||||
/**
|
||||
* Load Pyodide, install packages, and set up the Cognite SDK.
|
||||
* Safe to call multiple times — subsequent calls are no-ops.
|
||||
*/
|
||||
async initialize(sdk: PyodideSDKConfig): Promise<void> {
|
||||
if (this._initialized) return;
|
||||
|
||||
const report = this.config.onProgress ?? (() => {});
|
||||
const cdnUrl = this.config.cdnUrl ?? DEFAULT_CDN_URL;
|
||||
|
||||
// 1. Load Pyodide
|
||||
report('Loading Pyodide...', 10);
|
||||
this.pyodide = await this.config.loadPyodide({ indexURL: cdnUrl });
|
||||
report('Pyodide loaded', 30);
|
||||
|
||||
// 2. Core packages (micropip + HTTP patching)
|
||||
report('Loading core packages...', 40);
|
||||
await this.pyodide.loadPackage(['micropip', 'pyodide-http']);
|
||||
await this.pyodide.runPythonAsync(`
|
||||
try:
|
||||
import pyodide_http
|
||||
pyodide_http.patch_all()
|
||||
except Exception:
|
||||
pass
|
||||
`);
|
||||
this.micropip = this.pyodide.pyimport('micropip') as Micropip;
|
||||
|
||||
// 3. Cognite SDK
|
||||
const verb = isCacheValid() ? 'Loading' : 'Downloading';
|
||||
report(`${verb} cognite-sdk...`, 60);
|
||||
await this.micropip.install('cognite-sdk');
|
||||
if (!isCacheValid()) markCacheValid();
|
||||
report('cognite-sdk ready', 80);
|
||||
|
||||
// 4. Additional packages
|
||||
const reqs = this.config.requirements ?? [];
|
||||
if (reqs.length > 0) {
|
||||
report('Installing packages...', 85);
|
||||
await this.micropip.install(reqs);
|
||||
}
|
||||
|
||||
// 5. Utility functions + Cognite client
|
||||
report('Setting up environment...', 90);
|
||||
this.pyodide.runPython(PYTHON_UTILS);
|
||||
|
||||
report('Initializing Cognite client...', 95);
|
||||
this.pyodide.runPython(`
|
||||
import os
|
||||
os.environ["COGNITE_PROJECT"] = "${sdk.project}"
|
||||
os.environ["COGNITE_BASE_URL"] = "${sdk.baseUrl}"
|
||||
os.environ["COGNITE_TOKEN"] = "${sdk.token}"
|
||||
os.environ["COGNITE_FUSION_NOTEBOOK"] = "1"
|
||||
os.environ["MPLBACKEND"] = "AGG"
|
||||
from cognite.client import CogniteClient
|
||||
client = CogniteClient()
|
||||
`);
|
||||
|
||||
this._initialized = true;
|
||||
report('Ready', 100);
|
||||
}
|
||||
|
||||
/** Execute Python code asynchronously. PyProxy results are converted to JSON-safe values. */
|
||||
async runCodeAsync(code: string): Promise<unknown> {
|
||||
const pyodide = this.requirePyodide();
|
||||
const raw = await pyodide.runPythonAsync(code);
|
||||
return this.toJsonSafe(raw);
|
||||
}
|
||||
|
||||
/** Refresh the Cognite SDK token (e.g. after token rotation). */
|
||||
refreshToken(token: string): void {
|
||||
this.requirePyodide().runPython(
|
||||
`import os; os.environ["COGNITE_TOKEN"] = "${token}"`,
|
||||
);
|
||||
}
|
||||
|
||||
private requirePyodide(): PyodideInstance {
|
||||
if (!this.pyodide) {
|
||||
throw new Error(
|
||||
'PyodideRuntime not initialized — call initialize() first',
|
||||
);
|
||||
}
|
||||
return this.pyodide;
|
||||
}
|
||||
|
||||
/** Convert a Pyodide result to a JSON-safe JS value. */
|
||||
private toJsonSafe(value: unknown): unknown {
|
||||
if (value == null) return undefined;
|
||||
if (!isPyProxy(value)) return value;
|
||||
|
||||
const pyodide = this.requirePyodide();
|
||||
const converter = pyodide.globals.get('as_json_string') as
|
||||
| ((obj: unknown) => string)
|
||||
| undefined;
|
||||
|
||||
if (!converter) {
|
||||
throw new Error(
|
||||
'as_json_string not available — was initialize() called?',
|
||||
);
|
||||
}
|
||||
|
||||
try {
|
||||
return JSON.parse((converter as (obj: unknown) => string)(value));
|
||||
} finally {
|
||||
destroyIfPyProxy(converter);
|
||||
destroyIfPyProxy(value);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Singleton
|
||||
// ============================================================================
|
||||
|
||||
let globalInstance: PyodideRuntime | undefined;
|
||||
|
||||
/**
|
||||
* Get or create the global PyodideRuntime singleton.
|
||||
* Config is only used on first call — subsequent calls return the existing instance.
|
||||
*/
|
||||
export function getGlobalPyodideRuntime(
|
||||
config: PyodideRuntimeConfig,
|
||||
): PyodideRuntime {
|
||||
if (!globalInstance) {
|
||||
globalInstance = new PyodideRuntime(config);
|
||||
}
|
||||
return globalInstance;
|
||||
}
|
||||
|
||||
/** Reset the global runtime (e.g. on logout). */
|
||||
export function resetGlobalPyodideRuntime(): void {
|
||||
globalInstance = undefined;
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
/**
|
||||
* Pyodide entry point — browser-based Python execution for Atlas agents.
|
||||
*
|
||||
* Import path: @cognite/dune-utils/atlas-agent/pyodide
|
||||
*/
|
||||
|
||||
// Runtime
|
||||
export {
|
||||
PyodideRuntime,
|
||||
getGlobalPyodideRuntime,
|
||||
resetGlobalPyodideRuntime,
|
||||
clearPyodideCache,
|
||||
} from './pyodide-runtime';
|
||||
export type { PyodideRuntimeConfig, PyodideSDKConfig, PyodideInstance } from './pyodide-runtime';
|
||||
|
||||
// React hook
|
||||
export { usePyodideRuntime } from './pyodide-react';
|
||||
export type {
|
||||
PyodideProgress,
|
||||
UsePyodideRuntimeOptions,
|
||||
UsePyodideRuntimeReturn,
|
||||
} from './pyodide-react';
|
||||
@@ -0,0 +1,45 @@
|
||||
/**
|
||||
* Python execution primitives for Pyodide-based tool execution.
|
||||
*
|
||||
* buildWrapper — combines tool code + args into a single runnable Python snippet
|
||||
* formatOutput — serialises the Pyodide result into a string for the agent
|
||||
*/
|
||||
|
||||
/**
|
||||
* Unicode-safe base64 encoding.
|
||||
* Encodes to UTF-8 bytes first so non-ASCII characters survive round-tripping.
|
||||
*/
|
||||
function toBase64(str: string): string {
|
||||
const bytes = new TextEncoder().encode(str);
|
||||
let binary = '';
|
||||
for (const b of bytes) {
|
||||
binary += String.fromCharCode(b);
|
||||
}
|
||||
return btoa(binary);
|
||||
}
|
||||
|
||||
/**
|
||||
* Build a Python wrapper that loads the tool code, parses base64-encoded args,
|
||||
* and calls handle(**args). Supports both sync and async handle functions.
|
||||
*/
|
||||
export function buildWrapper(code: string, argsJson: string): string {
|
||||
const encoded = toBase64(argsJson);
|
||||
return `
|
||||
import json, base64, inspect
|
||||
_args = json.loads(base64.b64decode("${encoded}").decode("utf-8"))
|
||||
${code}
|
||||
async def _exec():
|
||||
if "handle" not in globals():
|
||||
return {"_error": "No handle() function found in tool code"}
|
||||
return await handle(**_args) if inspect.iscoroutinefunction(handle) else handle(**_args)
|
||||
_r = await _exec()
|
||||
json.dumps(_r) if _r is not None and not isinstance(_r, str) else _r
|
||||
`.trimStart();
|
||||
}
|
||||
|
||||
/** Stringify Python result into tool output text. */
|
||||
export function formatOutput(raw: unknown): string {
|
||||
if (raw == null) return '';
|
||||
if (typeof raw === 'string') return raw;
|
||||
return JSON.stringify(raw);
|
||||
}
|
||||
@@ -0,0 +1,278 @@
|
||||
/**
|
||||
* useAtlasChat — plug-and-play React hook for Atlas agent conversations.
|
||||
*
|
||||
* Manages session lifecycle, message state, streaming, and abort support.
|
||||
* Separate entry point from core for tree-shaking.
|
||||
*/
|
||||
|
||||
import { useState, useCallback, useRef, useEffect, useMemo } from 'react';
|
||||
import type { CogniteClient } from '@cognite/sdk';
|
||||
import { AtlasSession } from './session';
|
||||
import type { AtlasTool, AtlasResponse, PythonRuntime, ToolCall } from './types';
|
||||
|
||||
// ============================================================================
|
||||
// Types
|
||||
// ============================================================================
|
||||
|
||||
export interface ChatMessage<TContext = unknown> {
|
||||
id: string;
|
||||
role: 'user' | 'assistant';
|
||||
text: string;
|
||||
timestamp: Date;
|
||||
isStreaming?: boolean;
|
||||
/** Tool calls (client-side and server-side) attached to this message */
|
||||
toolCalls?: ToolCall[];
|
||||
/** App-specific context data, populated via onResponse */
|
||||
context?: TContext;
|
||||
}
|
||||
|
||||
export interface UseAtlasChatOptions<TContext = unknown> {
|
||||
client: CogniteClient | null;
|
||||
agentExternalId: string;
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||||
tools?: AtlasTool<any, any>[];
|
||||
/** Opt-in Python runtime (e.g. Pyodide) — only needed for agents that use Python tools. */
|
||||
pythonRuntime?: PythonRuntime;
|
||||
/** Messages to show on initial render (e.g. welcome message) */
|
||||
initialMessages?: ChatMessage<TContext>[];
|
||||
/** Called when a full response is received. Return context to merge into the assistant message. */
|
||||
onResponse?: (response: AtlasResponse) => TContext | void;
|
||||
/** Called before each send to inject app-level context (e.g. current todo state) into the request. */
|
||||
getAppContext?: () => string | undefined;
|
||||
}
|
||||
|
||||
export interface UseAtlasChatReturn<TContext = unknown> {
|
||||
/** All messages in the conversation */
|
||||
messages: ChatMessage<TContext>[];
|
||||
/** Send a user message — automatically creates user + assistant messages, handles streaming */
|
||||
send: (text: string) => Promise<void>;
|
||||
/** True while the agent is responding */
|
||||
isStreaming: boolean;
|
||||
/** Current progress text (e.g. "Agent thinking", "Executing: render_widget") */
|
||||
progress: string | null;
|
||||
/** Error message if last send failed */
|
||||
error: string | null;
|
||||
/** Clear all messages and reset the session */
|
||||
reset: () => void;
|
||||
/** Cancel the current streaming response */
|
||||
abort: () => void;
|
||||
/** Replace messages (e.g. loading conversation history) */
|
||||
setMessages: (messages: ChatMessage<TContext>[]) => void;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Helpers
|
||||
// ============================================================================
|
||||
|
||||
let messageCounter = 0;
|
||||
|
||||
function generateId(): string {
|
||||
return `msg-${Date.now()}-${++messageCounter}`;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Hook
|
||||
// ============================================================================
|
||||
|
||||
export function useAtlasChat<TContext = unknown>(
|
||||
options: UseAtlasChatOptions<TContext>,
|
||||
): UseAtlasChatReturn<TContext> {
|
||||
const { client, agentExternalId, tools, pythonRuntime, initialMessages, onResponse, getAppContext } = options;
|
||||
|
||||
const [messages, setMessages] = useState<ChatMessage<TContext>[]>(initialMessages ?? []);
|
||||
const [isStreaming, setIsStreaming] = useState(false);
|
||||
const [progress, setProgress] = useState<string | null>(null);
|
||||
const [error, setError] = useState<string | null>(null);
|
||||
|
||||
const sessionRef = useRef<AtlasSession | null>(null);
|
||||
const abortRef = useRef<AbortController | null>(null);
|
||||
const agentExternalIdRef = useRef(agentExternalId);
|
||||
const toolsRef = useRef(tools);
|
||||
const pythonRuntimeRef = useRef(pythonRuntime);
|
||||
const getAppContextRef = useRef(getAppContext);
|
||||
|
||||
// Keep refs updated (array/object identity may change between renders)
|
||||
toolsRef.current = tools;
|
||||
pythonRuntimeRef.current = pythonRuntime;
|
||||
getAppContextRef.current = getAppContext;
|
||||
|
||||
// Stable wrapper — always delegates to the latest getAppContext via ref.
|
||||
// Passed to AtlasSession once at creation so the session is never stale.
|
||||
const stableGetAppContext = useMemo(
|
||||
() => () => getAppContextRef.current?.(),
|
||||
[],
|
||||
);
|
||||
|
||||
// Cleanup on unmount
|
||||
useEffect(() => {
|
||||
return () => {
|
||||
abortRef.current?.abort();
|
||||
};
|
||||
}, []);
|
||||
|
||||
const getSession = useCallback((): AtlasSession | null => {
|
||||
if (!client) return null;
|
||||
|
||||
if (!sessionRef.current || agentExternalIdRef.current !== agentExternalId) {
|
||||
sessionRef.current = new AtlasSession({
|
||||
client,
|
||||
agentExternalId,
|
||||
tools: toolsRef.current,
|
||||
pythonRuntime: pythonRuntimeRef.current,
|
||||
getAppContext: stableGetAppContext,
|
||||
});
|
||||
agentExternalIdRef.current = agentExternalId;
|
||||
}
|
||||
|
||||
return sessionRef.current;
|
||||
}, [client, agentExternalId]);
|
||||
|
||||
const send = useCallback(
|
||||
async (text: string) => {
|
||||
const session = getSession();
|
||||
if (!session || isStreaming) return;
|
||||
|
||||
setError(null);
|
||||
setIsStreaming(true);
|
||||
setProgress('Agent thinking');
|
||||
|
||||
// Add user message
|
||||
const userMessage: ChatMessage<TContext> = {
|
||||
id: generateId(),
|
||||
role: 'user',
|
||||
text,
|
||||
timestamp: new Date(),
|
||||
};
|
||||
|
||||
const assistantId = generateId();
|
||||
let accumulatedText = '';
|
||||
let assistantCreated = false;
|
||||
|
||||
setMessages((prev) => [...prev, userMessage]);
|
||||
|
||||
const abortController = new AbortController();
|
||||
abortRef.current = abortController;
|
||||
|
||||
// ---- Helpers scoped to this send() call ----
|
||||
|
||||
/** Update a single message by id */
|
||||
const updateMsg = (id: string, updates: Partial<ChatMessage<TContext>>) => {
|
||||
setMessages((prev) =>
|
||||
prev.map((m) => (m.id === id ? { ...m, ...updates } : m)),
|
||||
);
|
||||
};
|
||||
|
||||
/** Finalize the assistant message — update if already created, otherwise add a new one */
|
||||
const finalizeAssistant = (fields: Partial<ChatMessage<TContext>>) => {
|
||||
if (assistantCreated) {
|
||||
updateMsg(assistantId, { isStreaming: false, ...fields });
|
||||
} else {
|
||||
setMessages((prev) => [
|
||||
...prev,
|
||||
{
|
||||
id: assistantId,
|
||||
role: 'assistant' as const,
|
||||
timestamp: new Date(),
|
||||
text: '',
|
||||
isStreaming: false,
|
||||
...fields,
|
||||
},
|
||||
]);
|
||||
}
|
||||
};
|
||||
|
||||
try {
|
||||
const response = await session.send(
|
||||
text,
|
||||
{
|
||||
onProgress: (progressText) => {
|
||||
setProgress(progressText);
|
||||
},
|
||||
onChunk: (chunk) => {
|
||||
if (!assistantCreated) {
|
||||
assistantCreated = true;
|
||||
setMessages((prev) => [
|
||||
...prev,
|
||||
{
|
||||
id: assistantId,
|
||||
role: 'assistant' as const,
|
||||
text: chunk,
|
||||
timestamp: new Date(),
|
||||
isStreaming: true,
|
||||
},
|
||||
]);
|
||||
}
|
||||
accumulatedText += chunk;
|
||||
updateMsg(assistantId, { text: accumulatedText });
|
||||
},
|
||||
onToolStart: (toolName) => {
|
||||
setProgress(`Executing: ${toolName}`);
|
||||
},
|
||||
},
|
||||
abortController.signal,
|
||||
);
|
||||
|
||||
// Finalize assistant message
|
||||
finalizeAssistant({
|
||||
text:
|
||||
response.text ||
|
||||
(assistantCreated
|
||||
? undefined
|
||||
: "I apologize, but I couldn't generate a response. Please try again."),
|
||||
toolCalls:
|
||||
response.toolCalls.length > 0
|
||||
? response.toolCalls
|
||||
: undefined,
|
||||
});
|
||||
|
||||
// Let the app attach context (e.g. applications) to the message
|
||||
const ctx = onResponse?.(response);
|
||||
if (ctx !== undefined) {
|
||||
updateMsg(assistantId, { context: ctx });
|
||||
}
|
||||
} catch (err) {
|
||||
if ((err as Error).name === 'AbortError') {
|
||||
// Cancelled by user — finalize any in-progress message
|
||||
if (assistantCreated) {
|
||||
updateMsg(assistantId, { isStreaming: false });
|
||||
}
|
||||
} else {
|
||||
const errorText =
|
||||
err instanceof Error ? err.message : 'Unknown error';
|
||||
setError(errorText);
|
||||
finalizeAssistant({ text: `Error: ${errorText}` });
|
||||
}
|
||||
} finally {
|
||||
setIsStreaming(false);
|
||||
setProgress(null);
|
||||
abortRef.current = null;
|
||||
}
|
||||
},
|
||||
[getSession, isStreaming, onResponse],
|
||||
);
|
||||
|
||||
const reset = useCallback(() => {
|
||||
abortRef.current?.abort();
|
||||
setMessages(initialMessages ?? []);
|
||||
setIsStreaming(false);
|
||||
setProgress(null);
|
||||
setError(null);
|
||||
sessionRef.current = null;
|
||||
}, [initialMessages]);
|
||||
|
||||
const abort = useCallback(() => {
|
||||
abortRef.current?.abort();
|
||||
}, []);
|
||||
|
||||
return {
|
||||
messages,
|
||||
send,
|
||||
isStreaming,
|
||||
progress,
|
||||
error,
|
||||
reset,
|
||||
abort,
|
||||
setMessages,
|
||||
};
|
||||
}
|
||||
|
||||
@@ -0,0 +1,155 @@
|
||||
import { describe, it, expect, vi } from 'vitest';
|
||||
import { AtlasSession } from './session';
|
||||
import type { AtlasSessionConfig, RawAgentResponse } from './types';
|
||||
|
||||
/**
|
||||
* Minimal mock that satisfies CogniteClient just enough for AtlasSession.
|
||||
* AtlasClient.post is the only call path we exercise, so we stub it via the
|
||||
* prototype after construction.
|
||||
*/
|
||||
function createMockConfig(
|
||||
overrides?: Partial<AtlasSessionConfig>,
|
||||
): AtlasSessionConfig {
|
||||
return {
|
||||
client: {} as AtlasSessionConfig['client'],
|
||||
agentExternalId: 'test-agent',
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
/** Build a raw response with a tool action that requests a client tool call. */
|
||||
function responseWithToolAction(
|
||||
actionId: string,
|
||||
toolName: string,
|
||||
args: Record<string, unknown>,
|
||||
): RawAgentResponse {
|
||||
return {
|
||||
agentId: 'test-agent',
|
||||
agentExternalId: 'test-agent',
|
||||
response: {
|
||||
type: 'result',
|
||||
cursor: 'cursor-1',
|
||||
messages: [
|
||||
{
|
||||
role: 'assistant',
|
||||
actions: [
|
||||
{
|
||||
type: 'clientTool',
|
||||
actionId,
|
||||
clientTool: { name: toolName, arguments: args },
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
/** Build a terminal response (no actions). */
|
||||
function terminalResponse(text: string): RawAgentResponse {
|
||||
return {
|
||||
agentId: 'test-agent',
|
||||
agentExternalId: 'test-agent',
|
||||
response: {
|
||||
type: 'result',
|
||||
cursor: 'cursor-2',
|
||||
messages: [{ role: 'assistant', content: { type: 'text', text } }],
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
describe(AtlasSession.name, () => {
|
||||
let postSpy: ReturnType<typeof vi.fn>;
|
||||
|
||||
function createSession(config?: Partial<AtlasSessionConfig>): AtlasSession {
|
||||
const session = new AtlasSession(createMockConfig(config));
|
||||
|
||||
// Stub the internal client.post so we never hit the network.
|
||||
postSpy = vi.fn();
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||||
(session as any).client = { post: postSpy };
|
||||
|
||||
return session;
|
||||
}
|
||||
|
||||
describe('appContext in continuation turns', () => {
|
||||
it('includes contextInformation on the initial user message', async () => {
|
||||
const session = createSession({
|
||||
getAppContext: () => 'todo state here',
|
||||
});
|
||||
|
||||
postSpy.mockResolvedValueOnce(terminalResponse('Done'));
|
||||
|
||||
await session.send('hello');
|
||||
|
||||
const payload = postSpy.mock.calls[0][0];
|
||||
expect(payload.contextInformation).toEqual({
|
||||
appContext: 'todo state here',
|
||||
});
|
||||
});
|
||||
|
||||
it('includes contextInformation on continuation turns after tool execution', async () => {
|
||||
let callCount = 0;
|
||||
const session = createSession({
|
||||
getAppContext: () => {
|
||||
callCount++;
|
||||
return `context-v${callCount}`;
|
||||
},
|
||||
tools: [
|
||||
{
|
||||
name: 'TestTool',
|
||||
description: 'test',
|
||||
parameters: { type: 'object', properties: {} },
|
||||
execute: () => ({ output: 'ok' }),
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
// Turn 1: agent requests a tool call
|
||||
postSpy.mockResolvedValueOnce(
|
||||
responseWithToolAction('action-1', 'TestTool', {}),
|
||||
);
|
||||
// Turn 2: terminal response
|
||||
postSpy.mockResolvedValueOnce(terminalResponse('All done'));
|
||||
|
||||
await session.send('do something');
|
||||
|
||||
// First call: initial payload
|
||||
expect(postSpy).toHaveBeenCalledTimes(2);
|
||||
const initialPayload = postSpy.mock.calls[0][0];
|
||||
expect(initialPayload.contextInformation).toEqual({
|
||||
appContext: 'context-v1',
|
||||
});
|
||||
|
||||
// Second call: continuation payload after tool execution
|
||||
const continuationPayload = postSpy.mock.calls[1][0];
|
||||
expect(continuationPayload.contextInformation).toEqual({
|
||||
appContext: 'context-v2',
|
||||
});
|
||||
});
|
||||
|
||||
it('omits contextInformation when getAppContext returns undefined', async () => {
|
||||
const session = createSession({
|
||||
getAppContext: () => undefined,
|
||||
});
|
||||
|
||||
postSpy.mockResolvedValueOnce(terminalResponse('Done'));
|
||||
|
||||
await session.send('hello');
|
||||
|
||||
const payload = postSpy.mock.calls[0][0];
|
||||
expect(payload.contextInformation).toBeUndefined();
|
||||
});
|
||||
|
||||
it('omits contextInformation when getAppContext is not provided', async () => {
|
||||
const session = createSession();
|
||||
|
||||
postSpy.mockResolvedValueOnce(terminalResponse('Done'));
|
||||
|
||||
await session.send('hello');
|
||||
|
||||
const payload = postSpy.mock.calls[0][0];
|
||||
expect(payload.contextInformation).toBeUndefined();
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,325 @@
|
||||
/**
|
||||
* AtlasSession — stateful conversation with validated tool execution.
|
||||
*
|
||||
* Manages the cursor, converts AtlasTool[] to API actions format,
|
||||
* validates tool arguments with ajv, and runs the tool execution loop.
|
||||
*/
|
||||
|
||||
import { AtlasClient } from './client';
|
||||
import { buildWrapper, formatOutput } from './python';
|
||||
import { validateToolArguments } from './validation';
|
||||
import type {
|
||||
AtlasTool,
|
||||
AtlasToolResult,
|
||||
AtlasResponse,
|
||||
AtlasSessionConfig,
|
||||
AgentToolConfig,
|
||||
PythonRuntime,
|
||||
StreamCallbacks,
|
||||
ChatPayload,
|
||||
ApiToolDefinition,
|
||||
RawAction,
|
||||
RawClientToolAction,
|
||||
RawAgentResponse,
|
||||
RequestMessage,
|
||||
ClientToolActionMessage,
|
||||
ToolCall,
|
||||
} from './types';
|
||||
|
||||
const MAX_TURNS = 50;
|
||||
|
||||
/**
|
||||
* Parse raw arguments from the API (string or object) into a plain object.
|
||||
*/
|
||||
function parseArguments(
|
||||
raw: string | Record<string, unknown>,
|
||||
): Record<string, unknown> {
|
||||
if (typeof raw === 'string') {
|
||||
try {
|
||||
return JSON.parse(raw);
|
||||
} catch {
|
||||
return {};
|
||||
}
|
||||
}
|
||||
return raw || {};
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract all actions from all messages in a raw response.
|
||||
*/
|
||||
function extractActions(raw: RawAgentResponse): RawAction[] {
|
||||
return raw.response.messages.flatMap((msg) => msg.actions ?? []);
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract server-side tool calls from reasoning blocks in a raw response.
|
||||
* Shape: messages[n].reasoning[n].data[n].toolCall
|
||||
*/
|
||||
function extractServerToolCalls(raw: RawAgentResponse): ToolCall[] {
|
||||
const calls: ToolCall[] = [];
|
||||
for (const msg of raw.response.messages) {
|
||||
for (const entry of msg.reasoning ?? []) {
|
||||
const data = (entry as { data?: unknown[] }).data;
|
||||
if (!Array.isArray(data)) continue;
|
||||
for (const item of data) {
|
||||
const tc = (item as { toolCall?: Record<string, unknown> }).toolCall;
|
||||
if (!tc) continue;
|
||||
const result = tc.result as Record<string, unknown> | undefined;
|
||||
calls.push({
|
||||
name: String(tc.name ?? ''),
|
||||
toolType: tc.toolType !== undefined ? String(tc.toolType) : undefined,
|
||||
input: tc.input,
|
||||
output: result?.output !== undefined ? String(result.output) : undefined,
|
||||
details: result?.result,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
return calls;
|
||||
}
|
||||
|
||||
/**
|
||||
* Execute a server-declared Python tool via the pythonRuntime.
|
||||
*/
|
||||
async function executePythonTool(
|
||||
action: RawClientToolAction,
|
||||
toolConfig: AgentToolConfig,
|
||||
pythonRuntime: PythonRuntime | undefined,
|
||||
callbacks?: StreamCallbacks,
|
||||
): Promise<{ result: AtlasToolResult; followup: ClientToolActionMessage }> {
|
||||
const toolName = action.clientTool.name;
|
||||
const pythonCode = String(toolConfig.configuration?.pythonCode ?? '').trim();
|
||||
|
||||
const fail = (msg: string) => {
|
||||
const result: AtlasToolResult = { output: msg };
|
||||
callbacks?.onToolEnd?.(toolName, result);
|
||||
return { result, followup: createActionReply(action.actionId, result.output) };
|
||||
};
|
||||
|
||||
if (!pythonCode) {
|
||||
return fail(`ERROR: pythonCode is empty in tool configuration for "${toolName}"`);
|
||||
}
|
||||
if (!pythonRuntime) {
|
||||
return fail(`ERROR: pythonRuntime is required to execute Python tool "${toolName}" but was not provided`);
|
||||
}
|
||||
|
||||
try {
|
||||
const argsJson = JSON.stringify(parseArguments(action.clientTool.arguments));
|
||||
const wrapper = buildWrapper(pythonCode, argsJson);
|
||||
const raw = await pythonRuntime.runCodeAsync(wrapper);
|
||||
const result: AtlasToolResult = { output: formatOutput(raw) };
|
||||
callbacks?.onToolEnd?.(toolName, result);
|
||||
return { result, followup: createActionReply(action.actionId, result.output) };
|
||||
} catch (err) {
|
||||
return fail(`ERROR: ${err instanceof Error ? err.message : String(err)}`);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Execute a client tool action: validate args, run the tool, return result + followup.
|
||||
*
|
||||
* Dispatch order:
|
||||
* 1. Pre-registered JS tool in the tools Map (client-declared tools)
|
||||
* 2. Server-defined Python tool — fetch config from agent API, run via pythonRuntime
|
||||
*/
|
||||
async function executeClientTool(
|
||||
action: RawClientToolAction,
|
||||
tools: Map<string, AtlasTool>,
|
||||
fetchToolConfig: (name: string) => Promise<AgentToolConfig | null>,
|
||||
pythonRuntime: PythonRuntime | undefined,
|
||||
callbacks?: StreamCallbacks,
|
||||
): Promise<{ result: AtlasToolResult; followup: ClientToolActionMessage }> {
|
||||
const toolName = action.clientTool.name;
|
||||
|
||||
callbacks?.onToolStart?.(toolName);
|
||||
|
||||
// 1. Pre-registered JS tool
|
||||
const tool = tools.get(toolName);
|
||||
if (tool) {
|
||||
const args = parseArguments(action.clientTool.arguments);
|
||||
try {
|
||||
validateToolArguments(toolName, tool.parameters, args);
|
||||
} catch (err) {
|
||||
const errorOutput = err instanceof Error ? err.message : String(err);
|
||||
const result: AtlasToolResult = { output: `ERROR: ${errorOutput}` };
|
||||
callbacks?.onToolEnd?.(toolName, result);
|
||||
return { result, followup: createActionReply(action.actionId, result.output) };
|
||||
}
|
||||
const result = await tool.execute(args);
|
||||
callbacks?.onToolEnd?.(toolName, result);
|
||||
return { result, followup: createActionReply(action.actionId, result.output) };
|
||||
}
|
||||
|
||||
// 2. Server-declared tool — look up config from agent API
|
||||
const toolConfig = await fetchToolConfig(toolName);
|
||||
if (toolConfig?.type === 'runPythonCode') {
|
||||
return executePythonTool(action, toolConfig, pythonRuntime, callbacks);
|
||||
}
|
||||
|
||||
const result: AtlasToolResult = { output: `Unknown client tool: ${toolName}` };
|
||||
callbacks?.onToolEnd?.(toolName, result);
|
||||
return { result, followup: createActionReply(action.actionId, result.output) };
|
||||
}
|
||||
|
||||
/**
|
||||
* Stateful conversation session with validated tool execution.
|
||||
*/
|
||||
export class AtlasSession {
|
||||
private cursor?: string;
|
||||
private readonly client: AtlasClient;
|
||||
private readonly agentExternalId: string;
|
||||
private readonly tools: Map<string, AtlasTool>;
|
||||
private readonly apiActionsOrUndefined: ApiToolDefinition[] | undefined;
|
||||
private readonly pythonRuntime: PythonRuntime | undefined;
|
||||
private readonly getAppContext: (() => string | undefined) | undefined;
|
||||
/** Cached tool configs fetched from the agent API (populated lazily on first Python tool call). */
|
||||
private cachedAgentTools: AgentToolConfig[] | undefined;
|
||||
|
||||
constructor(config: AtlasSessionConfig) {
|
||||
this.client = new AtlasClient(config.client);
|
||||
this.agentExternalId = config.agentExternalId;
|
||||
this.tools = new Map((config.tools || []).map((t) => [t.name, t]));
|
||||
this.pythonRuntime = config.pythonRuntime;
|
||||
this.getAppContext = config.getAppContext;
|
||||
|
||||
// Inline toApiToolDefinition (only used here)
|
||||
const apiActions: ApiToolDefinition[] = (config.tools || []).map((tool) => ({
|
||||
type: 'clientTool' as const,
|
||||
clientTool: {
|
||||
name: tool.name,
|
||||
description: tool.description,
|
||||
parameters: tool.parameters,
|
||||
},
|
||||
}));
|
||||
this.apiActionsOrUndefined = apiActions.length > 0 ? apiActions : undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
* Send a user message. Handles the full tool execution loop internally.
|
||||
*/
|
||||
async send(
|
||||
message: string,
|
||||
callbacks?: StreamCallbacks,
|
||||
signal?: AbortSignal,
|
||||
): Promise<AtlasResponse> {
|
||||
const allToolCalls: ToolCall[] = [];
|
||||
|
||||
const appContext = this.getAppContext?.();
|
||||
let payload: ChatPayload = {
|
||||
agentExternalId: this.agentExternalId,
|
||||
messages: [{ role: 'user', content: { type: 'text', text: message } }],
|
||||
actions: this.apiActionsOrUndefined,
|
||||
stream: true,
|
||||
...(this.cursor && { cursor: this.cursor }),
|
||||
...(appContext && { contextInformation: { appContext } }),
|
||||
};
|
||||
|
||||
for (let turn = 0; turn < MAX_TURNS; turn++) {
|
||||
const raw = await this.client.post(payload, this.agentExternalId, callbacks, signal);
|
||||
const response = raw.response;
|
||||
|
||||
if (response.type !== 'result') {
|
||||
throw new Error(`Unexpected response type: ${response.type}`);
|
||||
}
|
||||
|
||||
if (response.cursor) {
|
||||
this.cursor = response.cursor;
|
||||
}
|
||||
|
||||
// Collect server-side tool calls from reasoning blocks in this turn
|
||||
allToolCalls.push(...extractServerToolCalls(raw));
|
||||
|
||||
const actions = extractActions(raw);
|
||||
|
||||
// No actions → conversation turn is done
|
||||
if (actions.length === 0) {
|
||||
const text = response.messages?.[0]?.content?.text || '';
|
||||
return { text, cursor: this.cursor, toolCalls: allToolCalls, raw };
|
||||
}
|
||||
|
||||
// Execute actions and build follow-up messages
|
||||
const followups: RequestMessage[] = [];
|
||||
|
||||
for (const action of actions) {
|
||||
if (action.type === 'clientTool') {
|
||||
const { result, followup } = await executeClientTool(
|
||||
action,
|
||||
this.tools,
|
||||
(name) => this.fetchToolConfig(name),
|
||||
this.pythonRuntime,
|
||||
callbacks,
|
||||
);
|
||||
allToolCalls.push({ name: action.clientTool.name, output: result.output, details: result.details });
|
||||
followups.push(followup);
|
||||
} else if (action.type === 'toolConfirmation') {
|
||||
const toolName = action.toolConfirmation?.toolName;
|
||||
if (toolName) callbacks?.onToolStart?.(toolName);
|
||||
followups.push({
|
||||
role: 'action',
|
||||
type: 'toolConfirmation',
|
||||
actionId: action.actionId,
|
||||
status: 'ALLOW',
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
if (followups.length === 0) {
|
||||
const text = response.messages?.[0]?.content?.text || '';
|
||||
return { text, cursor: this.cursor, toolCalls: allToolCalls, raw };
|
||||
}
|
||||
|
||||
// Prepare the next turn — re-evaluate appContext so it reflects state changes from tool execution
|
||||
const updatedAppContext = this.getAppContext?.();
|
||||
payload = {
|
||||
agentExternalId: this.agentExternalId,
|
||||
messages: followups,
|
||||
actions: this.apiActionsOrUndefined,
|
||||
stream: true,
|
||||
cursor: this.cursor,
|
||||
...(updatedAppContext && { contextInformation: { appContext: updatedAppContext } }),
|
||||
};
|
||||
}
|
||||
|
||||
throw new Error(`Max tool execution turns reached (${MAX_TURNS})`);
|
||||
}
|
||||
|
||||
/**
|
||||
* Fetch a tool's config from the agent API (lazy, cached per session).
|
||||
* Used as a fallback when a clientTool action arrives for an unregistered tool.
|
||||
*/
|
||||
private async fetchToolConfig(toolName: string): Promise<AgentToolConfig | null> {
|
||||
if (!this.cachedAgentTools) {
|
||||
const agent = await this.client.getAgentById(this.agentExternalId);
|
||||
this.cachedAgentTools = agent?.tools ?? [];
|
||||
}
|
||||
return this.cachedAgentTools.find((t) => t.name === toolName) ?? null;
|
||||
}
|
||||
|
||||
/** Reset the session cursor (start a fresh conversation). */
|
||||
reset(): void {
|
||||
this.cursor = undefined;
|
||||
}
|
||||
|
||||
/** Get the current cursor value. */
|
||||
getCursor(): string | undefined {
|
||||
return this.cursor;
|
||||
}
|
||||
|
||||
/** Set the cursor (e.g. when restoring a conversation). */
|
||||
setCursor(cursor: string): void {
|
||||
this.cursor = cursor;
|
||||
}
|
||||
}
|
||||
|
||||
function createActionReply(
|
||||
actionId: string,
|
||||
text: string,
|
||||
): ClientToolActionMessage {
|
||||
return {
|
||||
role: 'action',
|
||||
type: 'clientTool',
|
||||
actionId,
|
||||
content: { type: 'text', text },
|
||||
data: [],
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,320 @@
|
||||
/**
|
||||
* Core types for the Atlas Agent client library.
|
||||
*
|
||||
* This module is self-contained — no imports from outside the library
|
||||
* except external packages (@sinclair/typebox, @cognite/sdk).
|
||||
*/
|
||||
|
||||
import type { TSchema, Static } from '@sinclair/typebox';
|
||||
import type { CogniteClient } from '@cognite/sdk';
|
||||
|
||||
// ============================================================================
|
||||
// Agent Types
|
||||
// ============================================================================
|
||||
|
||||
/** Configuration for a tool stored in the agent's CDF config. */
|
||||
export interface AgentToolConfig {
|
||||
name: string;
|
||||
type: string;
|
||||
configuration?: {
|
||||
pythonCode?: string;
|
||||
[key: string]: unknown;
|
||||
};
|
||||
[key: string]: unknown;
|
||||
}
|
||||
|
||||
export interface Agent {
|
||||
externalId: string;
|
||||
name: string;
|
||||
description?: string;
|
||||
model?: string;
|
||||
instructions?: string;
|
||||
ownerId?: string;
|
||||
tools?: AgentToolConfig[];
|
||||
createdTime?: number;
|
||||
lastUpdatedTime?: number;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Tool Types
|
||||
// ============================================================================
|
||||
|
||||
/** Result from executing a tool */
|
||||
export interface AtlasToolResult<TDetails = unknown> {
|
||||
/** Text sent back to the agent as tool output */
|
||||
output: string;
|
||||
/** Structured data for the app/UI */
|
||||
details?: TDetails;
|
||||
}
|
||||
|
||||
/**
|
||||
* A client-side tool the Atlas agent can invoke.
|
||||
* TypeBox schema for type-safe params + runtime validation via ajv.
|
||||
*/
|
||||
export interface AtlasTool<
|
||||
TParameters extends TSchema = TSchema,
|
||||
TDetails = unknown,
|
||||
> {
|
||||
name: string;
|
||||
description: string;
|
||||
parameters: TParameters;
|
||||
execute: (
|
||||
args: Static<TParameters>,
|
||||
) => AtlasToolResult<TDetails> | Promise<AtlasToolResult<TDetails>>;
|
||||
}
|
||||
|
||||
/** Minimal interface for executing Python code (e.g. Pyodide). */
|
||||
export interface PythonRuntime {
|
||||
runCodeAsync(code: string): Promise<unknown>;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Callback Types
|
||||
// ============================================================================
|
||||
|
||||
export interface StreamCallbacks {
|
||||
onProgress?: (text: string) => void;
|
||||
onChunk?: (text: string) => void;
|
||||
onToolStart?: (toolName: string) => void;
|
||||
onToolEnd?: (toolName: string, result: AtlasToolResult) => void;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Response Types (app-facing)
|
||||
// ============================================================================
|
||||
|
||||
/** A single tool invocation (client-side or server-side), ready for the UI. */
|
||||
export interface ToolCall {
|
||||
/** Friendly display name, e.g. "Find files" */
|
||||
name: string;
|
||||
/** Server-side tool type, e.g. "queryKnowledgeGraph" */
|
||||
toolType?: string;
|
||||
/** Raw input arguments */
|
||||
input?: unknown;
|
||||
/** Text returned to the agent as tool output */
|
||||
output?: string;
|
||||
/** Structured data for UI rendering */
|
||||
details?: unknown;
|
||||
}
|
||||
|
||||
export interface AtlasResponse {
|
||||
text: string;
|
||||
cursor?: string;
|
||||
toolCalls: ToolCall[];
|
||||
raw: RawAgentResponse;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Config Types
|
||||
// ============================================================================
|
||||
|
||||
export interface AtlasSessionConfig {
|
||||
client: CogniteClient;
|
||||
agentExternalId: string;
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||||
tools?: AtlasTool<any, any>[];
|
||||
/** Opt-in Python runtime (e.g. Pyodide) — required only for agents that use Python tools. */
|
||||
pythonRuntime?: PythonRuntime;
|
||||
/** Called before each send to inject app-level context (e.g. current todo state) into the request. */
|
||||
getAppContext?: () => string | undefined;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// API Shared Primitives (maps to cog_ai…session.common)
|
||||
// ============================================================================
|
||||
|
||||
/** Maps to AgentContentDTO */
|
||||
export interface AgentContent {
|
||||
type: string;
|
||||
text?: string;
|
||||
}
|
||||
|
||||
/** Maps to InstanceIdDTO */
|
||||
export interface InstanceId {
|
||||
space: string;
|
||||
externalId: string;
|
||||
}
|
||||
|
||||
/** Maps to ViewDTO */
|
||||
export interface View {
|
||||
space: string;
|
||||
externalId: string;
|
||||
version: string;
|
||||
}
|
||||
|
||||
/** Maps to PropertyVal type alias */
|
||||
export type PropertyVal =
|
||||
| string
|
||||
| number
|
||||
| boolean
|
||||
| null
|
||||
| Record<string, unknown>
|
||||
| (string | null)[]
|
||||
| (number | null)[]
|
||||
| (boolean | null)[]
|
||||
| (Record<string, unknown> | null)[];
|
||||
|
||||
// ============================================================================
|
||||
// Agent Data Types (maps to AgentDataDTO)
|
||||
// ============================================================================
|
||||
|
||||
/** Maps to InstanceDataDTO — the only variant we narrow on in app code. */
|
||||
export interface InstanceData {
|
||||
type: 'instance';
|
||||
view: View;
|
||||
instanceId: InstanceId;
|
||||
properties?: Record<string, PropertyVal>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Data items attached to response messages.
|
||||
* Only `InstanceData` is typed — other variants pass through as plain objects.
|
||||
*/
|
||||
export type AgentData = InstanceData | (Record<string, unknown> & { type: string });
|
||||
|
||||
// ============================================================================
|
||||
// Tool Definition Types (maps to CustomClientActionDTO)
|
||||
// ============================================================================
|
||||
|
||||
/** Maps to clientToolParameters */
|
||||
export interface ClientToolParameters {
|
||||
type: 'object';
|
||||
description?: string;
|
||||
properties?: Record<string, Record<string, unknown>>;
|
||||
required?: string[];
|
||||
propertyOrdering?: string[];
|
||||
}
|
||||
|
||||
/** Maps to CustomClientActionDTO */
|
||||
export interface ApiToolDefinition {
|
||||
type: 'clientTool';
|
||||
clientTool: {
|
||||
name: string;
|
||||
description?: string;
|
||||
parameters: ClientToolParameters;
|
||||
};
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Request Message Types (maps to RequestMessageDTO)
|
||||
// ============================================================================
|
||||
|
||||
/** Maps to AgentChatMessageUserRequestDTO */
|
||||
export interface UserRequestMessage {
|
||||
role: 'user';
|
||||
content: AgentContent;
|
||||
}
|
||||
|
||||
/** Maps to ClientToolCallActionMessageDTO */
|
||||
export interface ClientToolActionMessage {
|
||||
role: 'action';
|
||||
type: 'clientTool';
|
||||
actionId: string;
|
||||
content: AgentContent;
|
||||
data: Array<Record<string, unknown>>;
|
||||
}
|
||||
|
||||
/** Maps to UserConfirmationResponseDTO */
|
||||
export interface UserConfirmationMessage {
|
||||
role: 'action';
|
||||
type: 'toolConfirmation';
|
||||
actionId: string;
|
||||
status: 'ALLOW' | 'DENY';
|
||||
}
|
||||
|
||||
/** Maps to UserSessionResponseDTO */
|
||||
export interface UserSessionMessage {
|
||||
role: 'action';
|
||||
type: 'userSession';
|
||||
actionId: string;
|
||||
nonce: string;
|
||||
}
|
||||
|
||||
/** Maps to RequestMessageDTO (discriminated union) */
|
||||
export type RequestMessage =
|
||||
| UserRequestMessage
|
||||
| ClientToolActionMessage
|
||||
| UserConfirmationMessage
|
||||
| UserSessionMessage;
|
||||
|
||||
// ============================================================================
|
||||
// Session Context (maps to AgentSessionContextDTO)
|
||||
// ============================================================================
|
||||
|
||||
export interface AgentSessionContext {
|
||||
instanceSpaces?: string[];
|
||||
dataModels?: Array<Record<string, unknown>>;
|
||||
timeZone?: string;
|
||||
appContext?: string;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Chat Payload (maps to AgentSessionRequest)
|
||||
// ============================================================================
|
||||
|
||||
export interface ChatPayload {
|
||||
agentExternalId?: string;
|
||||
messages: RequestMessage[];
|
||||
actions?: ApiToolDefinition[];
|
||||
contextInformation?: AgentSessionContext;
|
||||
cursor?: string;
|
||||
stream: boolean;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Raw Response Types
|
||||
// ============================================================================
|
||||
|
||||
/** Response action: agent requests client to execute a tool */
|
||||
export interface RawClientToolAction {
|
||||
type: 'clientTool';
|
||||
actionId: string;
|
||||
clientTool: {
|
||||
name: string;
|
||||
arguments: string | Record<string, unknown>;
|
||||
};
|
||||
}
|
||||
|
||||
/** Response action: agent requests user confirmation */
|
||||
export interface RawToolConfirmationAction {
|
||||
type: 'toolConfirmation';
|
||||
actionId: string;
|
||||
toolConfirmation?: {
|
||||
toolName?: string;
|
||||
toolType?: string;
|
||||
toolArguments?: Record<string, unknown>;
|
||||
toolDescription?: string;
|
||||
content?: AgentContent;
|
||||
};
|
||||
}
|
||||
|
||||
/** Response action: agent requests user session */
|
||||
export interface RawUserSessionAction {
|
||||
type: 'userSession';
|
||||
actionId: string;
|
||||
}
|
||||
|
||||
/** Action from agent response (discriminated by `type`). Unknown types are skipped by the session loop. */
|
||||
export type RawAction =
|
||||
| RawClientToolAction
|
||||
| RawToolConfirmationAction
|
||||
| RawUserSessionAction;
|
||||
|
||||
/** A message in an agent response */
|
||||
export interface RawMessage {
|
||||
content?: AgentContent;
|
||||
role: string;
|
||||
data?: AgentData[];
|
||||
reasoning?: Array<Record<string, unknown>>;
|
||||
actions?: RawAction[];
|
||||
}
|
||||
|
||||
export interface RawAgentResponse {
|
||||
agentId: string;
|
||||
agentExternalId: string;
|
||||
response: {
|
||||
type: string;
|
||||
cursor?: string;
|
||||
messages: RawMessage[];
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
/**
|
||||
* Runtime validation for tool arguments using ajv.
|
||||
*
|
||||
* Ported from pi-mono packages/ai/src/utils/validation.ts pattern:
|
||||
* - Singleton ajv instance with coercion
|
||||
* - Graceful degradation if ajv fails to initialise (CSP)
|
||||
*/
|
||||
|
||||
import Ajv from 'ajv';
|
||||
import addFormats from 'ajv-formats';
|
||||
import type { TSchema } from '@sinclair/typebox';
|
||||
|
||||
let ajvInstance: Ajv | null = null;
|
||||
|
||||
function getAjv(): Ajv | null {
|
||||
if (ajvInstance) return ajvInstance;
|
||||
|
||||
try {
|
||||
ajvInstance = new Ajv({ allErrors: true, strict: false, coerceTypes: true });
|
||||
addFormats(ajvInstance);
|
||||
return ajvInstance;
|
||||
} catch {
|
||||
// Graceful degradation — skip validation if ajv cannot initialise
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Validate and coerce tool arguments against a TypeBox / JSON Schema.
|
||||
* Throws a formatted error on validation failure.
|
||||
* Gracefully skips validation when ajv is unavailable (e.g. CSP).
|
||||
*/
|
||||
export function validateToolArguments(
|
||||
toolName: string,
|
||||
schema: TSchema,
|
||||
args: unknown,
|
||||
): void {
|
||||
const ajv = getAjv();
|
||||
if (!ajv) return;
|
||||
|
||||
const validate = ajv.compile(schema);
|
||||
const valid = validate(args);
|
||||
if (valid) return;
|
||||
|
||||
const errors = validate.errors
|
||||
?.map((e) => `${e.instancePath || '/'} ${e.message}`)
|
||||
.join('; ');
|
||||
throw new Error(`Tool "${toolName}" received invalid arguments: ${errors}`);
|
||||
}
|
||||
Reference in New Issue
Block a user