Skip to content

Commit

Permalink
clean
Browse files Browse the repository at this point in the history
  • Loading branch information
@bmgalego
bmgalego committed Feb 22, 2025
1 parent 3c0ecde commit 6753645
Show file tree
Hide file tree
Showing 2 changed files with 82 additions and 79 deletions.
158 changes: 79 additions & 79 deletions packages/core/src/core/v1/dreams.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,11 +40,11 @@ export function createDreams<
any
>,
>(config: Config<Memory, TContext>): Agent<Memory, TContext> {
const taskRunner = new TaskRunner(3);
let booted = false;

const inputSubscriptions = new Map<string, Subscription>();

const contexts = new Map<string, { type: string; args?: any }>();
const contextsRunning = new Set<string>();

const {
inputs = {},
Expand All @@ -60,6 +60,8 @@ export function createDreams<

const container = config.container ?? createContainer();

const taskRunner = config.taskRunner ?? new TaskRunner(3);

const logger = new Logger({
level: config.logger ?? LogLevel.INFO,
enableTimestamp: true,
Expand All @@ -68,8 +70,6 @@ export function createDreams<

container.instance("logger", logger);

const contextsRunning = new Set<string>();

const debug: Debugger = (...args) => {
if (!config.debugger) return;
try {
Expand All @@ -79,10 +79,12 @@ export function createDreams<
}
};

let booted = false;

const serviceManager = createServiceManager(container);

for (const service of services) {
serviceManager.register(service);
}

for (const extension of extensions) {
if (extension.inputs) Object.assign(inputs, extension.inputs);
if (extension.outputs) Object.assign(outputs, extension.outputs);
Expand All @@ -91,61 +93,6 @@ export function createDreams<
if (extension.services) services.push(...extension.services);
}

for (const service of services) {
serviceManager.register(service);
}

function getContextId<TContext extends AnyContext>(
context: TContext,
args: z.infer<TContext["schema"]>
) {
const key = context.key ? context.key(args) : context.type;
return context.key ? [context.type, key].join(":") : context.type;
}

async function getContextState<TContext extends AnyContext>(
context: TContext,
args: z.infer<TContext["schema"]>
): Promise<ContextState<TContext>> {
const key = context.key ? context.key(args) : context.type;
const id = context.key ? [context.type, key].join(":") : context.type;

const options = context.setup ? await context.setup(args, agent) : {};

const memory =
(await agent.memory.store.get(id)) ??
(context.create
? context.create({ key, args, context, id: id, options })
: {});

return {
id,
key,
args,
options,
context,
memory,
};
}

async function getContextWorkingMemory(contextId: string) {
return (
(await agent.memory.store.get<WorkingMemory>(
["working-memory", contextId].join(":")
)) ?? (await defaultWorkingMemory.create())
);
}

async function saveContextWorkingMemory(
contextId: string,
workingMemory: WorkingMemory
) {
return await agent.memory.store.set(
["working-memory", contextId].join(":"),
workingMemory
);
}

const agent: Agent<Memory, TContext> = {
inputs,
outputs,
Expand All @@ -157,6 +104,7 @@ export function createDreams<
container,
model,
reasoningModel,
taskRunner,
debugger: debug,
context: config.context ?? undefined,
emit: (event: string, data: any) => {
Expand All @@ -172,15 +120,15 @@ export function createDreams<
},

getContext(params) {
return getContextState(params.context, params.args);
return getContextState(agent, params.context, params.args);
},

getContextId(params) {
return getContextId(params.context, params.args);
},

getWorkingMemory(contextId) {
return getContextWorkingMemory(contextId);
return getContextWorkingMemory(agent, contextId);
},

async start(args) {
Expand All @@ -199,11 +147,11 @@ export function createDreams<
if (input.install) await Promise.resolve(input.install(agent));

if (input.subscribe) {
let subscription = input.subscribe((contextHandler, args, data) => {
logger.info("agent", "input", { contextHandler, args, data });
let subscription = input.subscribe((context, args, data) => {
logger.info("agent", "input", { context, args, data });
agent
.send({
context: contextHandler,
context,
input: { type, data },
args,
})
Expand All @@ -229,7 +177,7 @@ export function createDreams<
}

if (agent.context) {
const { id } = await getContextState(agent.context, args);
const { id } = await getContextState(agent, agent.context, args);
contexts.set(id, { type: agent.context.type, args });
contexts.set("agent:context", { type: agent.context.type, args });
}
Expand All @@ -253,7 +201,7 @@ export function createDreams<
run: async ({ context, args, outputs, handlers }) => {
if (!booted) throw new Error("Not booted");

const ctxState = await getContextState(context, args);
const ctxState = await getContextState(agent, context, args);

contexts.set(ctxState.id, { type: context.type, args });

Expand All @@ -265,7 +213,7 @@ export function createDreams<
if (contextsRunning.has(ctxState.id)) return [];
contextsRunning.add(ctxState.id);

const workingMemory = await getContextWorkingMemory(ctxState.id);
const workingMemory = await getContextWorkingMemory(agent, ctxState.id);

const contextOuputs: Output[] = Object.entries({
...agent.outputs,
Expand Down Expand Up @@ -313,6 +261,7 @@ export function createDreams<

const agentCtxState = agent.context
? await getContextState(
agent,
agent.context,
contexts.get("agent:context")!.args
)
Expand Down Expand Up @@ -361,9 +310,8 @@ export function createDreams<

await handleStream(stream, state.index, handler);

const data = await response;

logger.debug("agent:parsed", "data", data);
// const data = await response;
// logger.debug("agent:parsed", "data", data);

await Promise.allSettled(actionCalls);

Expand All @@ -376,7 +324,7 @@ export function createDreams<
await agent.memory.store.set(agentCtxState.id, agentCtxState.memory);
}

await saveContextWorkingMemory(ctxState.id, workingMemory);
await saveContextWorkingMemory(agent, ctxState.id, workingMemory);

step++;

Expand All @@ -394,7 +342,7 @@ export function createDreams<

await agent.memory.store.set(ctxState.id, ctxState.memory);

await saveContextWorkingMemory(ctxState.id, workingMemory);
await saveContextWorkingMemory(agent, ctxState.id, workingMemory);

contextsRunning.delete(ctxState.id);

Expand All @@ -411,11 +359,12 @@ export function createDreams<
options,
memory,
} = await getContextState(
agent,
params.context,
params.context.schema.parse(params.args)
);

const workingMemory = await getContextWorkingMemory(contextId);
const workingMemory = await getContextWorkingMemory(agent, contextId);

const input = agent.inputs[params.input.type];
const data = input.schema.parse(params.input.data);
Expand Down Expand Up @@ -452,10 +401,8 @@ export function createDreams<

await agent.memory.store.set(contextId, memory);
await agent.memory.vector.upsert(contextId, [memory]);
await agent.memory.store.set(
["working-memory", contextId].join(":"),
workingMemory
);

await saveContextWorkingMemory(agent, contextId, workingMemory);

return await agent.run(params);
},
Expand All @@ -471,6 +418,59 @@ export function createDreams<
return agent;
}

function getContextId<TContext extends AnyContext>(
context: TContext,
args: z.infer<TContext["schema"]>
) {
const key = context.key ? context.key(args) : context.type;
return context.key ? [context.type, key].join(":") : context.type;
}

async function getContextState<TContext extends AnyContext>(
agent: AnyAgent,
context: TContext,
args: z.infer<TContext["schema"]>
): Promise<ContextState<TContext>> {
const key = context.key ? context.key(args) : context.type;
const id = context.key ? [context.type, key].join(":") : context.type;

const options = context.setup ? await context.setup(args, agent) : {};

const memory =
(await agent.memory.store.get(id)) ??
(context.create
? context.create({ key, args, context, id: id, options })
: {});

return {
id,
key,
args,
options,
context,
memory,
};
}

async function getContextWorkingMemory(agent: AnyAgent, contextId: string) {
return (
(await agent.memory.store.get<WorkingMemory>(
[contextId, "working-memory"].join(":")
)) ?? (await defaultWorkingMemory.create())
);
}

async function saveContextWorkingMemory(
agent: AnyAgent,
contextId: string,
workingMemory: WorkingMemory
) {
return await agent.memory.store.set(
[contextId, "working-memory"].join(":"),
workingMemory
);
}

const actionParseErrorPrompt = createPrompt(
`
You are tasked with fixing an action call arguments parsing error!
Expand Down
3 changes: 3 additions & 0 deletions packages/core/src/core/v1/types.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ import { z } from "zod";
import type { Container } from "./container";
import type { ServiceProvider } from "./serviceProvider";
import type { BaseMemory } from "./memory";
import type { TaskRunner } from "./task";

/**
* Represents a memory configuration for storing data
Expand Down Expand Up @@ -404,6 +405,8 @@ export interface Agent<

container: Container;

taskRunner: TaskRunner;

model: LanguageModelV1;
reasoningModel?: LanguageModelV1;

Expand Down

0 comments on commit 6753645

Please sign in to comment.