diff --git a/docs/concepts/agents.mdx b/docs/concepts/agents.mdx index 7f2c934..1fc3ca3 100644 --- a/docs/concepts/agents.mdx +++ b/docs/concepts/agents.mdx @@ -40,6 +40,43 @@ const agent = new Agent({ Read more about [Tools](/concepts/tools) to learn how to create them. +## Adding a retrieval function + +Agents can be configured with a retrieval function that will be used to retrieve information from the conversation history. This is useful for agents that need to reference previous messages in the conversation. + +One of the primary use cases for retrieval functions is to allow agents to reference an ontology of information that the developer wants the agent to have access to. Perhaps a company Wiki for example. This can typically be done by vectorizing the supplied message. + +```typescript +const retriever = async (messages: ChatCompletionMessageParam) => { + // Hit the embedding model to get the query embedding + const response = await openai.embeddings.create({ + model: "text-embedding-3-large", // Use the latest embedding model + input: messages + .filter(({ content }) => !!content) + .map(({ content }) => content).join("\n"), + }); + + const embedding = response.data[0].embedding; + + // Query the vector database with the query embedding + const { data: documents } = await supabase + .rpc('match_documents', { + query_embedding: embedding, + match_threshold: 0.8, + match_count: 10 + }); + + return documents.map(({ content }) => ({ content })); +}; + +const agent = new Agent({ + //... + retrievers: [retriever], +}); +``` + +The implementation of the `retriever` function has been inspired by the [guide by Supbase](https://supabase.com/blog/openai-embeddings-postgres-vector) on getting a vector database up and running. That said, the implementation is entirely up to you to decide. + ## Running an Agent Agents can be run standalone or as part of a ZEE workflow. Running an agent standalone is useful for testing and debugging. The `run` method will return the final state of the agent. diff --git a/packages/ai-agent-sdk/src/core/agent/index.ts b/packages/ai-agent-sdk/src/core/agent/index.ts index 0aec12e..cda3c8d 100644 --- a/packages/ai-agent-sdk/src/core/agent/index.ts +++ b/packages/ai-agent-sdk/src/core/agent/index.ts @@ -7,6 +7,8 @@ import type { ParsedFunctionToolCall } from "openai/resources/beta/chat/completi import type { ChatCompletionMessageParam } from "openai/resources/chat/completions"; import z from "zod"; +type Retrieval = { content: string }; + type AgentConfig = { name: string; model: ModelConfig; @@ -15,6 +17,13 @@ type AgentConfig = { instructions?: string[]; tools?: Record; + + retrievers?: (( + messages: ChatCompletionMessageParam[], + agent: Agent, + state: ZeeWorkflowState + ) => Promise)[]; + runFn?: ( agent: Agent, state: ZeeWorkflowState @@ -52,9 +61,34 @@ const defaultFn = async ( agent: Agent, state: ZeeWorkflowState ): Promise => { + const retrievers = agent.retrievers ?? []; + + const result = await Promise.all( + retrievers.map(async (retriever) => { + try { + return await retriever(state.messages, agent, state); + } catch (error) { + console.error("Error retrieving documents", error); + // Fail silently + return null; + } + }) + ); + + const retrievals = result.filter((r): r is Retrieval[] => r !== null); + const messages = [ system(` - ${agent.description} + ${agent.description}${retrievals + .map( + (retrieval) => ` + + + ${retrieval.map(({ content }) => content).join("\n")} + + ` + ) + .join("")} Your job is to complete the assigned task: - You can break down complex tasks into multiple steps if needed. @@ -286,6 +320,10 @@ export class Agent extends Base { return this.config.instructions; } + get retrievers() { + return this.config.retrievers; + } + get tools(): Record { return this._tools; }