mirror of
https://github.com/getcompanion-ai/co-mono.git
synced 2026-04-15 18:01:22 +00:00
Refactor OAuth/API key handling: AuthStorage and ModelRegistry
- Add AuthStorage class for credential storage (auth.json) - Add ModelRegistry class for model management with API key resolution - Add discoverAuthStorage() and discoverModels() discovery functions - Add migration from legacy oauth.json and settings.json apiKeys to auth.json - Remove configureOAuthStorage, defaultGetApiKey, findModel, discoverAvailableModels - Remove apiKeys from Settings type and SettingsManager methods - Rename getOAuthPath to getAuthPath - Update SDK, examples, docs, tests, and mom package Fixes #296
This commit is contained in:
Mario Zechner
parent
9f97f0c8da
commit
54018b6cc0
29 changed files with 953 additions and 2017 deletions
|
|
@ -13,8 +13,8 @@
|
|||
* pi --hook examples/hooks/custom-compaction.ts
|
||||
*/
|
||||
|
||||
import { complete } from "@mariozechner/pi-ai";
|
||||
import { findModel, messageTransformer } from "@mariozechner/pi-coding-agent";
|
||||
import { complete, getModel } from "@mariozechner/pi-ai";
|
||||
import { messageTransformer } from "@mariozechner/pi-coding-agent";
|
||||
import type { HookAPI } from "@mariozechner/pi-coding-agent/hooks";
|
||||
|
||||
export default function (pi: HookAPI) {
|
||||
|
|
@ -27,10 +27,9 @@ export default function (pi: HookAPI) {
|
|||
event;
|
||||
|
||||
// Use Gemini Flash for summarization (cheaper/faster than most conversation models)
|
||||
// findModel searches both built-in models and custom models from models.json
|
||||
const { model, error } = findModel("google", "gemini-2.5-flash");
|
||||
if (error || !model) {
|
||||
ctx.ui.notify(`Could not find Gemini Flash model: ${error}, using default compaction`, "warning");
|
||||
const model = getModel("google", "gemini-2.5-flash");
|
||||
if (!model) {
|
||||
ctx.ui.notify(`Could not find Gemini Flash model, using default compaction`, "warning");
|
||||
return;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -4,16 +4,27 @@
|
|||
* Shows how to select a specific model and thinking level.
|
||||
*/
|
||||
|
||||
import { createAgentSession, discoverAvailableModels, findModel } from "../../src/index.js";
|
||||
import { getModel } from "@mariozechner/pi-ai";
|
||||
import { createAgentSession, discoverAuthStorage, discoverModels } from "../../src/index.js";
|
||||
|
||||
// Option 1: Find a specific model by provider/id
|
||||
const { model: sonnet } = findModel("anthropic", "claude-sonnet-4-20250514");
|
||||
if (sonnet) {
|
||||
console.log(`Found model: ${sonnet.provider}/${sonnet.id}`);
|
||||
// Set up auth storage and model registry
|
||||
const authStorage = discoverAuthStorage();
|
||||
const modelRegistry = discoverModels(authStorage);
|
||||
|
||||
// Option 1: Find a specific built-in model by provider/id
|
||||
const opus = getModel("anthropic", "claude-opus-4-5");
|
||||
if (opus) {
|
||||
console.log(`Found model: ${opus.provider}/${opus.id}`);
|
||||
}
|
||||
|
||||
// Option 2: Pick from available models (have valid API keys)
|
||||
const available = await discoverAvailableModels();
|
||||
// Option 2: Find model via registry (includes custom models from models.json)
|
||||
const customModel = modelRegistry.find("my-provider", "my-model");
|
||||
if (customModel) {
|
||||
console.log(`Found custom model: ${customModel.provider}/${customModel.id}`);
|
||||
}
|
||||
|
||||
// Option 3: Pick from available models (have valid API keys)
|
||||
const available = await modelRegistry.getAvailable();
|
||||
console.log(
|
||||
"Available models:",
|
||||
available.map((m) => `${m.provider}/${m.id}`),
|
||||
|
|
@ -23,6 +34,8 @@ if (available.length > 0) {
|
|||
const { session } = await createAgentSession({
|
||||
model: available[0],
|
||||
thinkingLevel: "medium", // off, low, medium, high
|
||||
authStorage,
|
||||
modelRegistry,
|
||||
});
|
||||
|
||||
session.subscribe((event) => {
|
||||
|
|
|
|||
|
|
@ -1,40 +1,55 @@
|
|||
/**
|
||||
* API Keys and OAuth
|
||||
*
|
||||
* Configure API key resolution. Default checks: models.json, OAuth, env vars.
|
||||
* Configure API key resolution via AuthStorage and ModelRegistry.
|
||||
*/
|
||||
|
||||
import { getAgentDir } from "../../src/config.js";
|
||||
import { configureOAuthStorage, createAgentSession, defaultGetApiKey, SessionManager } from "../../src/index.js";
|
||||
import {
|
||||
AuthStorage,
|
||||
createAgentSession,
|
||||
discoverAuthStorage,
|
||||
discoverModels,
|
||||
ModelRegistry,
|
||||
SessionManager,
|
||||
} from "../../src/index.js";
|
||||
|
||||
// Default: uses env vars (ANTHROPIC_API_KEY, etc.), OAuth, and models.json
|
||||
await createAgentSession({
|
||||
sessionManager: SessionManager.inMemory(),
|
||||
});
|
||||
console.log("Session with default API key resolution");
|
||||
|
||||
// Custom resolver
|
||||
await createAgentSession({
|
||||
getApiKey: async (model) => {
|
||||
// Custom logic (secrets manager, database, etc.)
|
||||
if (model.provider === "anthropic") {
|
||||
return process.env.MY_ANTHROPIC_KEY;
|
||||
}
|
||||
// Fall back to default
|
||||
return defaultGetApiKey()(model);
|
||||
},
|
||||
sessionManager: SessionManager.inMemory(),
|
||||
});
|
||||
console.log("Session with custom API key resolver");
|
||||
|
||||
// Use OAuth from ~/.pi/agent while customizing everything else
|
||||
configureOAuthStorage(getAgentDir()); // Must call before createAgentSession
|
||||
// Default: discoverAuthStorage() uses ~/.pi/agent/auth.json
|
||||
// discoverModels() loads built-in + custom models from ~/.pi/agent/models.json
|
||||
const authStorage = discoverAuthStorage();
|
||||
const modelRegistry = discoverModels(authStorage);
|
||||
|
||||
await createAgentSession({
|
||||
agentDir: "/tmp/custom-config", // Custom config location
|
||||
// But OAuth tokens still come from ~/.pi/agent/oauth.json
|
||||
systemPrompt: "You are helpful.",
|
||||
skills: [],
|
||||
sessionManager: SessionManager.inMemory(),
|
||||
authStorage,
|
||||
modelRegistry,
|
||||
});
|
||||
console.log("Session with OAuth from default location, custom config elsewhere");
|
||||
console.log("Session with default auth storage and model registry");
|
||||
|
||||
// Custom auth storage location
|
||||
const customAuthStorage = new AuthStorage("/tmp/my-app/auth.json");
|
||||
const customModelRegistry = new ModelRegistry(customAuthStorage, "/tmp/my-app/models.json");
|
||||
|
||||
await createAgentSession({
|
||||
sessionManager: SessionManager.inMemory(),
|
||||
authStorage: customAuthStorage,
|
||||
modelRegistry: customModelRegistry,
|
||||
});
|
||||
console.log("Session with custom auth storage location");
|
||||
|
||||
// Runtime API key override (not persisted to disk)
|
||||
authStorage.setRuntimeApiKey("anthropic", "sk-my-temp-key");
|
||||
await createAgentSession({
|
||||
sessionManager: SessionManager.inMemory(),
|
||||
authStorage,
|
||||
modelRegistry,
|
||||
});
|
||||
console.log("Session with runtime API key override");
|
||||
|
||||
// No models.json - only built-in models
|
||||
const simpleRegistry = new ModelRegistry(authStorage); // null = no models.json
|
||||
await createAgentSession({
|
||||
sessionManager: SessionManager.inMemory(),
|
||||
authStorage,
|
||||
modelRegistry: simpleRegistry,
|
||||
});
|
||||
console.log("Session with only built-in models");
|
||||
|
|
|
|||
|
|
@ -2,38 +2,36 @@
|
|||
* Full Control
|
||||
*
|
||||
* Replace everything - no discovery, explicit configuration.
|
||||
* Still uses OAuth from ~/.pi/agent for convenience.
|
||||
*
|
||||
* IMPORTANT: When providing `tools` with a custom `cwd`, use the tool factory
|
||||
* functions (createReadTool, createBashTool, etc.) to ensure tools resolve
|
||||
* paths relative to your cwd.
|
||||
*/
|
||||
|
||||
import { getModel } from "@mariozechner/pi-ai";
|
||||
import { Type } from "@sinclair/typebox";
|
||||
import { getAgentDir } from "../../src/config.js";
|
||||
import {
|
||||
AuthStorage,
|
||||
type CustomAgentTool,
|
||||
configureOAuthStorage,
|
||||
createAgentSession,
|
||||
createBashTool,
|
||||
createReadTool,
|
||||
defaultGetApiKey,
|
||||
findModel,
|
||||
type HookFactory,
|
||||
ModelRegistry,
|
||||
SessionManager,
|
||||
SettingsManager,
|
||||
} from "../../src/index.js";
|
||||
|
||||
// Use OAuth from default location
|
||||
configureOAuthStorage(getAgentDir());
|
||||
// Custom auth storage location
|
||||
const authStorage = new AuthStorage("/tmp/my-agent/auth.json");
|
||||
|
||||
// Custom API key with fallback
|
||||
const getApiKey = async (model: { provider: string }) => {
|
||||
if (model.provider === "anthropic" && process.env.MY_ANTHROPIC_KEY) {
|
||||
return process.env.MY_ANTHROPIC_KEY;
|
||||
}
|
||||
return defaultGetApiKey()(model as any);
|
||||
};
|
||||
// Runtime API key override (not persisted)
|
||||
if (process.env.MY_ANTHROPIC_KEY) {
|
||||
authStorage.setRuntimeApiKey("anthropic", process.env.MY_ANTHROPIC_KEY);
|
||||
}
|
||||
|
||||
// Model registry with no custom models.json
|
||||
const modelRegistry = new ModelRegistry(authStorage);
|
||||
|
||||
// Inline hook
|
||||
const auditHook: HookFactory = (api) => {
|
||||
|
|
@ -55,7 +53,7 @@ const statusTool: CustomAgentTool = {
|
|||
}),
|
||||
};
|
||||
|
||||
const { model } = findModel("anthropic", "claude-sonnet-4-20250514");
|
||||
const model = getModel("anthropic", "claude-opus-4-5");
|
||||
if (!model) throw new Error("Model not found");
|
||||
|
||||
// In-memory settings with overrides
|
||||
|
|
@ -73,7 +71,8 @@ const { session } = await createAgentSession({
|
|||
|
||||
model,
|
||||
thinkingLevel: "off",
|
||||
getApiKey,
|
||||
authStorage,
|
||||
modelRegistry,
|
||||
|
||||
systemPrompt: `You are a minimal assistant.
|
||||
Available: read, bash, status. Be concise.`,
|
||||
|
|
|
|||
|
|
@ -29,50 +29,63 @@ npx tsx examples/sdk/01-minimal.ts
|
|||
## Quick Reference
|
||||
|
||||
```typescript
|
||||
import { getModel } from "@mariozechner/pi-ai";
|
||||
import {
|
||||
AuthStorage,
|
||||
createAgentSession,
|
||||
configureOAuthStorage,
|
||||
discoverAuthStorage,
|
||||
discoverModels,
|
||||
discoverSkills,
|
||||
discoverHooks,
|
||||
discoverCustomTools,
|
||||
discoverContextFiles,
|
||||
discoverSlashCommands,
|
||||
discoverAvailableModels,
|
||||
findModel,
|
||||
defaultGetApiKey,
|
||||
loadSettings,
|
||||
buildSystemPrompt,
|
||||
ModelRegistry,
|
||||
SessionManager,
|
||||
codingTools,
|
||||
readOnlyTools,
|
||||
readTool, bashTool, editTool, writeTool,
|
||||
} from "@mariozechner/pi-coding-agent";
|
||||
|
||||
// Auth and models setup
|
||||
const authStorage = discoverAuthStorage();
|
||||
const modelRegistry = discoverModels(authStorage);
|
||||
|
||||
// Minimal
|
||||
const { session } = await createAgentSession();
|
||||
const { session } = await createAgentSession({ authStorage, modelRegistry });
|
||||
|
||||
// Custom model
|
||||
const { model } = findModel("anthropic", "claude-sonnet-4-20250514");
|
||||
const { session } = await createAgentSession({ model, thinkingLevel: "high" });
|
||||
const model = getModel("anthropic", "claude-opus-4-5");
|
||||
const { session } = await createAgentSession({ model, thinkingLevel: "high", authStorage, modelRegistry });
|
||||
|
||||
// Modify prompt
|
||||
const { session } = await createAgentSession({
|
||||
systemPrompt: (defaultPrompt) => defaultPrompt + "\n\nBe concise.",
|
||||
authStorage,
|
||||
modelRegistry,
|
||||
});
|
||||
|
||||
// Read-only
|
||||
const { session } = await createAgentSession({ tools: readOnlyTools });
|
||||
const { session } = await createAgentSession({ tools: readOnlyTools, authStorage, modelRegistry });
|
||||
|
||||
// In-memory
|
||||
const { session } = await createAgentSession({
|
||||
sessionManager: SessionManager.inMemory(),
|
||||
authStorage,
|
||||
modelRegistry,
|
||||
});
|
||||
|
||||
// Full control
|
||||
configureOAuthStorage(); // Use OAuth from ~/.pi/agent
|
||||
const customAuth = new AuthStorage("/my/app/auth.json");
|
||||
customAuth.setRuntimeApiKey("anthropic", process.env.MY_KEY!);
|
||||
const customRegistry = new ModelRegistry(customAuth);
|
||||
|
||||
const { session } = await createAgentSession({
|
||||
model,
|
||||
getApiKey: async (m) => process.env.MY_KEY,
|
||||
authStorage: customAuth,
|
||||
modelRegistry: customRegistry,
|
||||
systemPrompt: "You are helpful.",
|
||||
tools: [readTool, bashTool],
|
||||
customTools: [{ tool: myTool }],
|
||||
|
|
@ -81,7 +94,6 @@ const { session } = await createAgentSession({
|
|||
contextFiles: [],
|
||||
slashCommands: [],
|
||||
sessionManager: SessionManager.inMemory(),
|
||||
settings: { compaction: { enabled: false } },
|
||||
});
|
||||
|
||||
// Run prompts
|
||||
|
|
@ -97,11 +109,12 @@ await session.prompt("Hello");
|
|||
|
||||
| Option | Default | Description |
|
||||
|--------|---------|-------------|
|
||||
| `authStorage` | `discoverAuthStorage()` | Credential storage |
|
||||
| `modelRegistry` | `discoverModels(authStorage)` | Model registry |
|
||||
| `cwd` | `process.cwd()` | Working directory |
|
||||
| `agentDir` | `~/.pi/agent` | Config directory |
|
||||
| `model` | From settings/first available | Model to use |
|
||||
| `thinkingLevel` | From settings/"off" | off, low, medium, high |
|
||||
| `getApiKey` | Built-in resolver | API key function |
|
||||
| `systemPrompt` | Discovered | String or `(default) => modified` |
|
||||
| `tools` | `codingTools` | Built-in tools |
|
||||
| `customTools` | Discovered | Replaces discovery |
|
||||
|
|
@ -112,7 +125,7 @@ await session.prompt("Hello");
|
|||
| `contextFiles` | Discovered | AGENTS.md files |
|
||||
| `slashCommands` | Discovered | File commands |
|
||||
| `sessionManager` | `SessionManager.create(cwd)` | Persistence |
|
||||
| `settings` | From agentDir | Overrides |
|
||||
| `settingsManager` | From agentDir | Settings overrides |
|
||||
|
||||
## Events
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue