getcompanion-ai/co-mono
1
0
Fork 0
mirror of https://github.com/getcompanion-ai/co-mono.git synced 2026-04-21 08:02:11 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Massive refactor of API

Browse source 
- Switch to function based API
- Anthropic SDK style async generator
- Fully typed with escape hatches for custom models
This commit is contained in:
Mario Zechner 2025-09-02 23:59:36 +02:00
parent 004de3c9d0
commit 66cefb236e
29 changed files with 5835 additions and 6225 deletions
Download patch file Download diff file Expand all files Collapse all files

41
packages/ai/src/models.ts
View file

@ -1,44 +1,39 @@
import { PROVIDERS } from "./models.generated.js";
import type { KnownProvider, Model, Usage } from "./types.js";
import type { Api, KnownProvider, Model, Usage } from "./types.js";
// Re-export Model type
export type { KnownProvider, Model } from "./types.js";
// Dynamic model registry initialized from PROVIDERS
const modelRegistry: Map<string, Map<string, Model>> = new Map();
const modelRegistry: Map<string, Map<string, Model<Api>>> = new Map();
// Initialize registry from PROVIDERS on module load
for (const [provider, models] of Object.entries(PROVIDERS)) {
const providerModels = new Map<string, Model>();
const providerModels = new Map<string, Model<Api>>();
for (const [id, model] of Object.entries(models)) {
providerModels.set(id, model as Model);
providerModels.set(id, model as Model<Api>);
}
modelRegistry.set(provider, providerModels);
}
/**
* Get a model from the registry - typed overload for known providers
*/
export function getModel<P extends KnownProvider>(provider: P, modelId: keyof (typeof PROVIDERS)[P]): Model;
export function getModel(provider: string, modelId: string): Model | undefined;
export function getModel(provider: any, modelId: any): Model | undefined {
return modelRegistry.get(provider)?.get(modelId);
type ModelApi<
TProvider extends KnownProvider,
TModelId extends keyof (typeof PROVIDERS)[TProvider],
> = (typeof PROVIDERS)[TProvider][TModelId] extends { api: infer TApi } ? (TApi extends Api ? TApi : never) : never;
export function getModel<TProvider extends KnownProvider, TModelId extends keyof (typeof PROVIDERS)[TProvider]>(
provider: TProvider,
modelId: TModelId,
): Model<ModelApi<TProvider, TModelId>>;
export function getModel<TApi extends Api>(provider: string, modelId: string): Model<TApi> | undefined;
export function getModel<TApi extends Api>(provider: any, modelId: any): Model<TApi> | undefined {
return modelRegistry.get(provider)?.get(modelId) as Model<TApi> | undefined;
}
/**
* Register a custom model
*/
export function registerModel(model: Model): void {
export function registerModel<TApi extends Api>(model: Model<TApi>): void {
if (!modelRegistry.has(model.provider)) {
modelRegistry.set(model.provider, new Map());
}
modelRegistry.get(model.provider)!.set(model.id, model);
}
/**
* Calculate cost for token usage
*/
export function calculateCost(model: Model, usage: Usage): Usage["cost"] {
export function calculateCost<TApi extends Api>(model: Model<TApi>, usage: Usage): Usage["cost"] {
usage.cost.input = (model.cost.input / 1000000) * usage.input;
usage.cost.output = (model.cost.output / 1000000) * usage.output;
usage.cost.cacheRead = (model.cost.cacheRead / 1000000) * usage.cacheRead;