rag provider

Author: alextitonis

packages/plugin-raggraph/src/index.ts

@@ -1,117 +1,12 @@
-import { elizaLogger } from "@ai16z/eliza";
-import {
-    Action,
-    HandlerCallback,
-    IAgentRuntime,
-    Memory,
-    Plugin,
-    State,
-} from "@ai16z/eliza";
-import { createGraphRAG } from "./driver";
-import { validateRaggraphConfig } from "./environment";
-
-export async function initializeRAGGraph(runtime: IAgentRuntime) {
-    const config = await validateRaggraphConfig(runtime);
-    return createGraphRAG({
-        neo4jUri: config.NEO4J_URI,
-        neo4jUser: config.NEO4J_USER,
-        neo4jPassword: config.NEO4J_PASSWORD,
-    });
-}
-
-const ragQuery: Action = {
-    name: "RAG_QUERY",
-    similes: [
-        "SEARCH_KNOWLEDGE",
-        "QUERY_GRAPH",
-        "FIND_RELATED",
-        "SEARCH_DOCUMENTS",
-        "RETRIEVE_INFO",
-    ],
-    description: "Query the knowledge graph for relevant information",
-    validate: async (runtime: IAgentRuntime, _message: Memory) => {
-        try {
-            await validateRaggraphConfig(runtime);
-            return true;
-        } catch {
-            return false;
-        }
-    },
-    handler: async (
-        runtime: IAgentRuntime,
-        message: Memory,
-        state: State,
-        _options: Record<string, any>,
-        callback: HandlerCallback
-    ) => {
-        elizaLogger.log("Composing state for message:", message);
-        state = (await runtime.composeState(message)) as State;
-
-        try {
-            const rag = await initializeRAGGraph(runtime);
-            const query = message.content.text;
-
-            elizaLogger.log("Executing RAG query:", query);
-            const result = await rag.query(query);
-
-            // Close the connection after query
-            await rag.close();
-
-            callback({
-                text: "Here's what I found:",
-                attachments: [
-                    {
-                        id: crypto.randomUUID(),
-                        title: "Search Results",
-                        source: "raggraph",
-                        description: "Knowledge Graph Query Results",
-                        text: result.fullContext,
-                        metadata: {
-                            confidence: result.confidence,
-                            sources: result.sources,
-                        },
-                    },
-                ],
-            });
-        } catch (error) {
-            elizaLogger.error("RAG query failed:", error);
-            throw error;
-        }
-    },
-    examples: [
-        [
-            {
-                user: "{{user1}}",
-                content: { text: "What do you know about neural networks?" },
-            },
-            {
-                user: "{{agentName}}",
-                content: {
-                    text: "Let me search our knowledge base for information about neural networks.",
-                    action: "RAG_QUERY",
-                },
-            },
-        ],
-        [
-            {
-                user: "{{user1}}",
-                content: { text: "Find information about machine learning" },
-            },
-            {
-                user: "{{agentName}}",
-                content: {
-                    text: "I'll search our knowledge graph for machine learning information.",
-                    action: "RAG_QUERY",
-                },
-            },
-        ],
-    ],
-} as Action;
+import { Plugin } from "@ai16z/eliza";
+import { ragGraphProvider } from "./provider";
 
 export const raggraphPlugin: Plugin = {
     name: "raggraph",
-    description: "Knowledge graph querying and retrieval",
-    actions: [ragQuery],
+    description: "RAGGraph Plugin for Eliza",
+    actions: [],
     evaluators: [],
-    providers: [],
+    providers: [ragGraphProvider],
 };
+
+export default raggraphPlugin;

packages/plugin-raggraph/src/provider.ts

@@ -0,0 +1,102 @@
+import {
+    IAgentRuntime,
+    ICacheManager,
+    Memory,
+    Provider,
+    State,
+} from "@ai16z/eliza";
+import { createGraphRAG } from "./driver";
+import { GraphRAGResponse } from "./types";
+import NodeCache from "node-cache";
+import * as path from "path";
+
+const PROVIDER_CONFIG = {
+    MAX_RETRIES: 3,
+    RETRY_DELAY: 2000,
+};
+
+export class RAGGraphProvider {
+    private cache: NodeCache;
+    private cacheKey: string = "raggraph";
+    private graphRAG;
+
+    constructor(
+        private neo4jUri: string,
+        private neo4jUser: string,
+        private neo4jPassword: string,
+        private cacheManager: ICacheManager
+    ) {
+        this.cache = new NodeCache({ stdTTL: 300 }); // Cache TTL set to 5 minutes
+        this.graphRAG = createGraphRAG({
+            neo4jUri: this.neo4jUri,
+            neo4jUser: this.neo4jUser,
+            neo4jPassword: this.neo4jPassword,
+        });
+    }
+
+    private async readFromCache<T>(key: string): Promise<T | null> {
+        const cached = await this.cacheManager.get<T>(
+            path.join(this.cacheKey, key)
+        );
+        return cached;
+    }
+
+    private async writeToCache<T>(key: string, data: T): Promise<void> {
+        await this.cacheManager.set(path.join(this.cacheKey, key), data, {
+            expires: Date.now() + 5 * 60 * 1000,
+        });
+    }
+
+    async query(userQuery: string): Promise<GraphRAGResponse> {
+        const cacheKey = `query-${userQuery}`;
+
+        // Try to get from cache first
+        const cached = await this.readFromCache<GraphRAGResponse>(cacheKey);
+        if (cached) {
+            return cached;
+        }
+
+        // If not in cache, perform the query
+        const response = await this.graphRAG.query(userQuery);
+
+        // Cache the result
+        await this.writeToCache(cacheKey, response);
+
+        return response;
+    }
+
+    async close(): Promise<void> {
+        await this.graphRAG.close();
+    }
+}
+
+const ragGraphProvider: Provider = {
+    get: async (
+        runtime: IAgentRuntime,
+        message: Memory,
+        _state?: State
+    ): Promise<string | null> => {
+        try {
+            const neo4jUri = runtime.getSetting("NEO4J_URI");
+            const neo4jUser = runtime.getSetting("NEO4J_USER");
+            const neo4jPassword = runtime.getSetting("NEO4J_PASSWORD");
+
+            const provider = new RAGGraphProvider(
+                neo4jUri,
+                neo4jUser,
+                neo4jPassword,
+                runtime.cacheManager
+            );
+
+            const response = await provider.query(message.content);
+            await provider.close();
+
+            return response.fullContext;
+        } catch (error) {
+            console.error("Error in RAG graph provider:", error);
+            return null;
+        }
+    },
+};
+
+export { ragGraphProvider };