FROM feat/332-add-memories TO development (#337)

ryaneggz · web-flow · commit 86df6ed9f015 · 2025-06-15T22:31:28.000-05:00
* Memories appera to work neew frontend toggle

* Can toggle memory

* Few fixes

* Most of fixes g2g

* Should be better than before

* Mmoery working

* Should fix frontend
diff --git a/Changelog.md b/Changelog.md
@@ -34,6 +34,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   - bugfix/28-cannot-auth-tools-list (2024-11-30)
 
 ### Changed 
+  - feat/332-add-memories (2025-06-15)
   - feat/290-api-as-a-tool-v2 (2025-06-08)
   - feat/327-agent-select-from-chatinput (2025-06-06)
   - feat/322-add-react-voice-viz (2025-05-29)
diff --git a/backend/src/routes/v0/llm/__init__.py b/backend/src/routes/v0/llm/__init__.py
@@ -99,14 +99,17 @@ async def new_thread(
     db: AsyncSession = Depends(get_async_db)
 ):
     try:
-        store = None
-        if hasattr(body, 'memory') and body.memory:
-            store = get_store_db()
-            await store.setup()
-        
-        controller = AgentController(db=db, user_id=user.id if user else None)
-        output_type = request.headers.get("accept", "application/json")
-        return await controller.query_thread(output_type=output_type, thread=body, store=store)
+        async with get_store_db() as store:
+            args = {
+                "output_type": request.headers.get("accept", "application/json"),
+                "thread": body
+            }
+            if hasattr(body, 'memory') and body.memory:
+                await store.setup()
+                args["store"] = store
+            
+            controller = AgentController(db=db, user_id=user.id if user else None)
+            return await controller.query_thread(**args)
     except httpx.HTTPStatusError as e:
         logger.error(f"Error creating new thread: {str(e)}")
         raise HTTPException(status_code=e.response.status_code , detail=str(e))
@@ -149,9 +152,17 @@ async def existing_thread(
 ):
     try:
         async with get_store_db() as store:
+            args = {
+                "output_type": request.headers.get("accept", "application/json"),
+                "thread": body,
+                "thread_id": thread_id
+            }
+            if hasattr(body, 'memory') and body.memory:
+                await store.setup()
+                args["store"] = store
+                
             controller = AgentController(db=db, user_id=user.id if user else None)
-            output_type = request.headers.get("accept", "application/json")
-            return await controller.query_thread(output_type=output_type, thread=body, store=store)
+            return await controller.query_thread(**args)
     except httpx.HTTPStatusError as e:
         logger.error(f"Error creating new thread: {str(e)}")
         raise HTTPException(status_code=e.response.status_code , detail=str(e))
diff --git a/backend/src/services/db.py b/backend/src/services/db.py
@@ -7,6 +7,7 @@
 from psycopg_pool import ConnectionPool
 from src.constants import DB_URI, CONNECTION_POOL_KWARGS
 from langgraph.store.postgres import AsyncPostgresStore
+from langchain.embeddings.base import init_embeddings
 
 MAX_CONNECTION_POOL_SIZE = None
 
@@ -33,7 +34,14 @@ def get_checkpoint_db() -> AsyncPostgresSaver:
     return AsyncPostgresSaver.from_conn_string(DB_URI)
 
 def get_store_db() -> AsyncPostgresStore:
-    return AsyncPostgresStore.from_conn_string(DB_URI)
+    return AsyncPostgresStore.from_conn_string(
+        conn_string=DB_URI,
+        index={
+            "dims": 1536,
+            "embed": init_embeddings("openai:text-embedding-3-small"),
+            "fields": ["memory"]  # specify which fields to embed. Default is the whole serialized value
+        }
+    )
 
 # Session context managers
 def get_db() -> Generator[SessionLocal, None, None]: # type: ignore
diff --git a/backend/src/tools/api.py b/backend/src/tools/api.py
@@ -184,37 +184,30 @@ def make_api_call_func(
         data: Optional request body data for POST/PUT requests
         path_params: Optional dictionary of path parameters to format into the URL
     """
-    async def api_call(event: str, data: Dict[str, Any], config: RunnableConfig):
+    async def api_call(event: str, data: Dict[str, Any] = {}):
         wrapper = GenericRequestsWrapper(headers=headers)
         # Format URL with path parameters if provided
         formatted_url = url.format(**(path_params or {}))
-        async with get_store_db() as store:
-            try:
-                if method == "GET":
-                    response = await wrapper.aget(formatted_url)
-                elif method == "POST":
-                    response = await wrapper.apost(formatted_url, data={'data': data, "event": event})
-                elif method == "PUT":
-                    response = await wrapper.aput(formatted_url, data={'data': data, "event": event})
-                elif method == "PATCH":
-                    response = await wrapper.apatch(formatted_url, data={'data': data, "event": event})
-                elif method == "DELETE":
-                    response = await wrapper.adelete(formatted_url)
-                else:
-                    raise ValueError(f"Unsupported HTTP method: {method}")
-                
-                logger.debug(f"API Response for {method} {formatted_url}: {response}")
-                if store and config.get("store", None):
-                    await store.aput(("memories", get_user_id(config)), event, {"memory": {
-                    "event": event,
-                    "data": data,
-                    "response": response
-                }})  
-                return response
-                
-            except Exception as e:
-                logger.error(f"API call failed for {method} {formatted_url}: {str(e)}")
-                raise
+        try:
+            if method == "GET":
+                response = await wrapper.aget(formatted_url)
+            elif method == "POST":
+                response = await wrapper.apost(formatted_url, data={'data': data, "event": event})
+            elif method == "PUT":
+                response = await wrapper.aput(formatted_url, data={'data': data, "event": event})
+            elif method == "PATCH":
+                response = await wrapper.apatch(formatted_url, data={'data': data, "event": event})
+            elif method == "DELETE":
+                response = await wrapper.adelete(formatted_url)
+            else:
+                raise ValueError(f"Unsupported HTTP method: {method}")
+            
+            logger.debug(f"API Response for {method} {formatted_url}: {response}")
+            return response
+            
+        except Exception as e:
+            logger.error(f"API call failed for {method} {formatted_url}: {str(e)}")
+            raise
 
     api_call.__doc__ = description
     return api_call
diff --git a/backend/src/tools/memory.py b/backend/src/tools/memory.py
@@ -0,0 +1,60 @@
+import uuid
+from typing import List, TypedDict
+from langchain_core.tools import tool, ToolException
+from langchain_core.documents import Document
+from langchain_core.runnables import RunnableConfig
+from src.services.db import get_store_db
+from src.utils.tools import get_user_id
+from src.utils.logger import logger
+
+class Memory(TypedDict):
+	memory: str
+	ttl: int = 10_080 # 1 week
+
+@tool
+async def save_recall_memory(
+    memory: str, 
+    ttl: int = 1440, 
+    config: RunnableConfig = None
+) -> str:
+	"""Save memory to vectorstore for later semantic retrieval."""
+	memory_id = str(uuid.uuid4())
+	async with get_store_db() as store:
+		await store.aput(("memories", get_user_id(config)), memory_id, {"memory": memory}, ttl=ttl)
+	return f"Memory ID {memory_id} saved."
+
+@tool
+async def delete_recall_memory(memory_id: str, config: RunnableConfig) -> str:
+	"""Delete a specific memory for the current thread."""
+	
+	async with get_store_db() as store:
+		await store.adelete(("memories", get_user_id(config)), memory_id)
+	return f"Memory ID {memory_id} deleted."
+
+## TODO: Not sure if this works correctly, does not appear that way.
+@tool()
+async def search_recall_memories(query: str, config: RunnableConfig) -> List[str]:
+	"""
+	Search for relevant memories.
+
+	Args:
+		query: The question to search for.
+
+	Returns:
+		A list of memories that are relevant to the query.
+	"""
+	try:
+		user_id = get_user_id(config)
+
+		async with get_store_db() as store:
+			documents = await store.asearch(
+				("memories", get_user_id(config)), 
+				query=query, 
+				limit=3, 
+				filter={"user_id": user_id}
+			)
+		print(documents)
+		return [doc.dict() for doc in documents]
+	except ToolException as e:
+		logger.exception(f"Error searching for memories: {e}")
+		raise ToolException(f"Error searching for memories: {str(e)}")
diff --git a/backend/src/utils/agent.py b/backend/src/utils/agent.py
@@ -1,5 +1,4 @@
 import asyncio
-import json
 from fastapi import HTTPException, status
 from fastapi.responses import Response, JSONResponse, StreamingResponse
 from langchain_core.messages import AnyMessage
@@ -13,7 +12,8 @@
 from src.schemas.entities.a2a import A2AServer
 from src.repos.user_repo import UserRepo
 from src.constants import APP_LOG_LEVEL
-from src.tools import dynamic_tools, init_tools
+from src.tools import init_tools
+from src.tools.memory import search_recall_memories, save_recall_memory, delete_recall_memory
 from src.utils.llm import LLMWrapper
 from src.constants.llm import ModelName
 from src.schemas.entities import Answer, Thread, ArcadeConfig
@@ -125,13 +125,13 @@ async def abuilder(
         )
         self.llm = LLMWrapper(model_name=model_name, tools=self.tools, user_repo=self.user_repo)
         if self.store:
-            self.store.embeddings = LLMWrapper(model_name="openai:text-embedding-3-large").embedding_model()
+            self.tools.extend([search_recall_memories, save_recall_memory, delete_recall_memory])
             memories = await self.store.asearch(("memories", self.user_id))
             if memories:
-                system += "\n\nMemories:\n"
+                system += "\n\n"
                 for memory in memories:
                     memory_dict = str(memory.dict())
-                    system += f"{memory_dict}\n"
+                    system += f"<recall_memory>{memory_dict}</recall_memory>\n"
                 
         if self.tools:
             graph = create_react_agent(self.llm, prompt=system, tools=self.tools, checkpointer=self.checkpointer, store=self.store)
diff --git a/backend/src/utils/tools.py b/backend/src/utils/tools.py
@@ -58,4 +58,11 @@ def get_user_id(config: RunnableConfig) -> str:
     if user_id is None:
         raise ValueError("User ID needs to be provided to save a memory.")
 
-    return user_id
+    return user_id
+
+def get_thread_id(config: RunnableConfig) -> str:
+    thread_id = config["configurable"].get("thread_id")
+    if thread_id is None:
+        raise ValueError("Thread ID needs to be provided to save a memory.")
+
+    return thread_id
diff --git a/frontend/src/components/drawers/ConfigDrawer/tabs/tab-content-info.tsx b/frontend/src/components/drawers/ConfigDrawer/tabs/tab-content-info.tsx
@@ -2,20 +2,58 @@ import { Button } from "@/components/ui/button"
 import SelectModel from "@/components/selects/SelectModel"
 import { useChatContext } from "@/context/ChatContext"
 import SystemMessageCard from "@/components/cards/SystemMessageCard"
+import { Switch } from "@/components/ui/switch"
+import { useMemory } from "@/hooks/useMemory"
+import { useSystem } from "@/hooks/useSystem"
+import { useState } from "react"
 
 function TabContentInfo() {
-	const { payload } = useChatContext()
+	const { payload, setPayload } = useChatContext()
+	const [completed, ] = useState(false)
+
+	const handleMemoryToggle = () => {
+		setPayload((prev: any) => ({ ...prev, memory: !prev.memory }));
+	}
+
+	const handleSaveSystemPrompt = () => {
+		setPayload((prev: any) => ({ ...prev, system: payload.system }));
+	}
+
   return (
     <div className="space-y-6">
 			<div className="space-y-4">
 				<h3 className="text-lg font-medium">Model</h3>
 				<SelectModel />
 			</div>
 
+			<div className="space-y-4">
+				<h3 className="text-lg font-medium">Memory</h3>
+				<div className="flex items-center justify-between">
+					<div className="space-y-1">
+						<p className="text-sm font-medium">Enable Memory</p>
+						<p className="text-xs text-muted-foreground">
+							Allow the AI to remember previous conversations
+						</p>
+					</div>
+					<Switch
+						checked={useMemory()}
+						onCheckedChange={handleMemoryToggle}
+					/>
+				</div>
+			</div>
+
 			<div className="space-y-4">
 				<h3 className="text-lg font-medium">System Prompt</h3>
-				<SystemMessageCard content={payload.system} />
-				<Button className="mt-4 w-full">Save</Button>
+				<SystemMessageCard content={useSystem()} />
+				<Button 
+					className="mt-4 w-full" 
+					onClick={() => {
+						handleSaveSystemPrompt();
+						alert("System prompt saved");
+					}}
+				>
+					{completed ? "Saved!" : "Save"}
+				</Button>
 			</div>
 		</div>
   )
diff --git a/frontend/src/components/inputs/ChatInput.tsx b/frontend/src/components/inputs/ChatInput.tsx
@@ -126,7 +126,7 @@ export default function ChatInput() {
         onDrop={handleDrop}
         onDragOver={(e) => e.preventDefault()}
         onKeyDown={(e) => {
-          if (e.key === "Enter" && !e.shiftKey) {
+          if (e.key === "Enter" && !e.shiftKey && !isRecording && payload.query.length > 0) {
             e.preventDefault()
             handleSubmit()
           }
diff --git a/frontend/src/components/selects/SelectModel/SelectModel.tsx b/frontend/src/components/selects/SelectModel/SelectModel.tsx
@@ -6,14 +6,14 @@ import { SiAnthropic, SiOpenai, SiOllama, SiGoogle } from 'react-icons/si';
 import GroqIcon from "@/components/icons/GroqIcon";
 import { useQueryParam, StringParam, withDefault } from 'use-query-params';
 import { useChatContext } from "@/context/ChatContext";
-import { DEFAULT_CHAT_MODEL } from "@/lib/config/llm";
+import { useModel } from "@/hooks/useModel";
+
 
-// Create a parameter with default value
-const ModelParam = withDefault(StringParam, DEFAULT_CHAT_MODEL);
 
 function SelectModel() {
+	// Create a parameter with default value
 	const { setPayload, models } = useChatContext();
-	const [model, setModel] = useQueryParam('model', ModelParam);
+	const [model, setModel] = useQueryParam('model', withDefault(StringParam, useModel()));
 
 	const handleModelChange = (modelId: string) => {
 		setModel(modelId);
@@ -24,7 +24,7 @@ function SelectModel() {
 		if (model) {
 			setPayload((prev: any) => ({ ...prev, model }));
 		}
-	}, [model, setPayload]);
+	}, [model]);
 
 	return (
 		<Select value={model} onValueChange={handleModelChange}>
diff --git a/frontend/src/hooks/useChatHook.ts b/frontend/src/hooks/useChatHook.ts
@@ -5,11 +5,11 @@ import apiClient from '@/lib/utils/apiClient';
 import { listModels, Model } from '@/lib/services/modelService';
 import { listTools } from '@/lib/services/toolService';
 import { useChatReducer } from '@/lib/reducers/chatReducer';
-import { DEFAULT_SYSTEM_PROMPT } from '@/lib/config/instruction';
 import { DEFAULT_CHAT_MODEL, isValidModelName } from '@/lib/config/llm';
 import { getAuthToken } from '@/lib/utils/auth';
 import { streamThread } from '@/lib/services/threadService';
 import { useAppContext } from '@/context/AppContext';
+import { getMemory, getModel, getSystemPrompt } from '@/lib/utils/storage';
 
 const KEY_NAME = 'config:mcp';
 
@@ -30,10 +30,11 @@ const initChatState = {
         threadId: '',
         images: [] as string[],
         query: '',
-        system: DEFAULT_SYSTEM_PROMPT,
+        system: getSystemPrompt(),
         tools: [] as any[],
         visualize: false,
-        model: '',
+        model: getModel(),
+        memory: getMemory(),
         mcp: null,
         a2a: null,
         collection: null,
diff --git a/frontend/src/hooks/useMemory.tsx b/frontend/src/hooks/useMemory.tsx
@@ -0,0 +1,26 @@
+import { useChatContext } from "@/context/ChatContext"
+import { useEffect } from "react"
+
+export function useMemory() {
+  const { payload, setPayload } = useChatContext();
+
+  useEffect(() => {
+    // Initialize memory from localStorage on mount
+    const storedMemory = localStorage.getItem("enso:chat:payload:memory");
+    if (storedMemory) {
+      const memory = JSON.parse(storedMemory);
+      if (typeof memory === 'boolean') {
+        setPayload((prev: any) => ({ ...prev, memory }));
+      }
+    }
+  }, []);
+
+  useEffect(() => {
+    // Update localStorage whenever payload.memory changes
+    if (typeof payload.memory === 'boolean') {
+      localStorage.setItem("enso:chat:payload:memory", JSON.stringify(payload.memory));
+    }
+  }, [payload.memory]);
+
+  return typeof payload.memory === 'boolean' ? payload.memory : false;
+}
diff --git a/frontend/src/hooks/useModel.tsx b/frontend/src/hooks/useModel.tsx
diff --git a/frontend/src/hooks/useSystem.tsx b/frontend/src/hooks/useSystem.tsx
diff --git a/frontend/src/lib/config/instruction.ts b/frontend/src/lib/config/instruction.ts
diff --git a/frontend/src/lib/entities/index.ts b/frontend/src/lib/entities/index.ts
diff --git a/frontend/src/lib/utils/storage.ts b/frontend/src/lib/utils/storage.ts

Original file line number	Diff line number	Diff line change
`@@ -126,7 +126,7 @@ export default function ChatInput() {`
`126`	`126`	`onDrop={handleDrop}`
`127`	`127`	`onDragOver={(e) => e.preventDefault()}`
`128`	`128`	`onKeyDown={(e) => {`
`129`		`- if (e.key === "Enter" && !e.shiftKey) {`
	`129`	`+ if (e.key === "Enter" && !e.shiftKey && !isRecording && payload.query.length > 0) {`
`130`	`130`	`e.preventDefault()`
`131`	`131`	`handleSubmit()`
`132`	`132`	`}`