WIP: Retrieval logic

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.

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<AgentName, Tool>;
+
+    retrievers?: ((
+        messages: ChatCompletionMessageParam[],
+        agent: Agent,
+        state: ZeeWorkflowState
+    ) => Promise<Retrieval[]>)[];
+
     runFn?: (
         agent: Agent,
         state: ZeeWorkflowState
@@ -52,9 +61,34 @@ const defaultFn = async (
     agent: Agent,
     state: ZeeWorkflowState
 ): Promise<ZeeWorkflowState> => {
+    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>
+                ${retrieval.map(({ content }) => content).join("\n")}
+            </retrieval>
+            `
+                )
+                .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<AgentName, Tool> {
         return this._tools;
     }