|
| 1 | +from lfx.base.memory.model import LCChatMemoryComponent |
| 2 | +from lfx.inputs.inputs import IntInput, MessageTextInput, SecretStrInput |
| 3 | +from lfx.io import Output |
| 4 | +from lfx.log.logger import logger |
| 5 | +from lfx.schema.data import Data |
| 6 | + |
| 7 | + |
| 8 | +class MengramMemoryComponent(LCChatMemoryComponent): |
| 9 | + display_name = "Mengram Memory" |
| 10 | + description = ( |
| 11 | + "Store and retrieve memories using Mengram — AI memory with semantic, episodic, and procedural types. " |
| 12 | + "Automatically extracts facts, events, and workflows from conversations." |
| 13 | + ) |
| 14 | + name = "mengram_memory" |
| 15 | + icon: str = "Mengram" |
| 16 | + |
| 17 | + inputs = [ |
| 18 | + SecretStrInput( |
| 19 | + name="mengram_api_key", |
| 20 | + display_name="Mengram API Key", |
| 21 | + info="API key for Mengram (starts with 'om-'). Get one at mengram.io.", |
| 22 | + ), |
| 23 | + MessageTextInput( |
| 24 | + name="user_id", |
| 25 | + display_name="User ID", |
| 26 | + info="Identifier for the user associated with the memories.", |
| 27 | + ), |
| 28 | + MessageTextInput( |
| 29 | + name="ingest_message", |
| 30 | + display_name="Message to Ingest", |
| 31 | + info="The message content to be ingested into Mengram memory.", |
| 32 | + ), |
| 33 | + MessageTextInput( |
| 34 | + name="search_query", |
| 35 | + display_name="Search Query", |
| 36 | + info="Input text for searching related memories across all memory types.", |
| 37 | + ), |
| 38 | + MessageTextInput( |
| 39 | + name="agent_id", |
| 40 | + display_name="Agent ID", |
| 41 | + info="Optional identifier for the agent sending the message.", |
| 42 | + advanced=True, |
| 43 | + ), |
| 44 | + MessageTextInput( |
| 45 | + name="app_id", |
| 46 | + display_name="App ID", |
| 47 | + info="Optional identifier for the application.", |
| 48 | + advanced=True, |
| 49 | + ), |
| 50 | + IntInput( |
| 51 | + name="top_k", |
| 52 | + display_name="Top K", |
| 53 | + info="Maximum number of results per memory type.", |
| 54 | + value=5, |
| 55 | + advanced=True, |
| 56 | + ), |
| 57 | + MessageTextInput( |
| 58 | + name="api_url", |
| 59 | + display_name="API URL", |
| 60 | + info="Mengram API base URL.", |
| 61 | + value="https://mengram.io", |
| 62 | + advanced=True, |
| 63 | + ), |
| 64 | + ] |
| 65 | + |
| 66 | + outputs = [ |
| 67 | + Output(name="search_results", display_name="Search Results", method="search_memories"), |
| 68 | + Output(name="cognitive_profile", display_name="Cognitive Profile", method="get_cognitive_profile"), |
| 69 | + Output(name="memory", display_name="Mengram Memory", method="ingest_data"), |
| 70 | + ] |
| 71 | + |
| 72 | + def _build_client(self): |
| 73 | + """Initialize a Mengram client instance.""" |
| 74 | + try: |
| 75 | + from mengram import Mengram |
| 76 | + except ImportError as e: |
| 77 | + msg = "Mengram is not installed. Please install it with 'pip install mengram-ai'." |
| 78 | + raise ImportError(msg) from e |
| 79 | + return Mengram(api_key=self.mengram_api_key, base_url=self.api_url) |
| 80 | + |
| 81 | + def ingest_data(self) -> Data: |
| 82 | + """Ingest a message into Mengram memory and return the result.""" |
| 83 | + client = self._build_client() |
| 84 | + |
| 85 | + if not self.ingest_message or not self.user_id: |
| 86 | + logger.warning("Missing 'ingest_message' or 'user_id'; cannot ingest data.") |
| 87 | + return Data(data={"status": "skipped", "reason": "missing ingest_message or user_id"}) |
| 88 | + |
| 89 | + logger.info("Ingesting message for user_id: %s", self.user_id) |
| 90 | + |
| 91 | + try: |
| 92 | + result = client.add( |
| 93 | + [{"role": "user", "content": self.ingest_message}], |
| 94 | + user_id=self.user_id, |
| 95 | + agent_id=self.agent_id or None, |
| 96 | + app_id=self.app_id or None, |
| 97 | + ) |
| 98 | + except Exception: |
| 99 | + logger.exception("Failed to add message to Mengram memory.") |
| 100 | + raise |
| 101 | + |
| 102 | + return Data(data={"status": "success", "result": result}) |
| 103 | + |
| 104 | + def search_memories(self) -> Data: |
| 105 | + """Search Mengram memory for related memories across all memory types.""" |
| 106 | + client = self._build_client() |
| 107 | + |
| 108 | + if not self.search_query or not self.user_id: |
| 109 | + logger.warning("Missing 'search_query' or 'user_id'; cannot search.") |
| 110 | + return Data(data={}) |
| 111 | + |
| 112 | + logger.info("Searching memories for user_id: %s", self.user_id) |
| 113 | + |
| 114 | + try: |
| 115 | + results = client.search_all( |
| 116 | + self.search_query, |
| 117 | + limit=self.top_k, |
| 118 | + user_id=self.user_id, |
| 119 | + ) |
| 120 | + except Exception: |
| 121 | + logger.exception("Failed to search Mengram memory.") |
| 122 | + raise |
| 123 | + |
| 124 | + return Data(data=results) |
| 125 | + |
| 126 | + def get_cognitive_profile(self) -> Data: |
| 127 | + """Get the Cognitive Profile — a system prompt generated from all memory types.""" |
| 128 | + client = self._build_client() |
| 129 | + |
| 130 | + if not self.user_id: |
| 131 | + logger.warning("Missing 'user_id'; cannot get cognitive profile.") |
| 132 | + return Data(data={"system_prompt": ""}) |
| 133 | + |
| 134 | + logger.info("Getting cognitive profile for user_id: %s", self.user_id) |
| 135 | + |
| 136 | + try: |
| 137 | + profile = client.get_profile(self.user_id) |
| 138 | + except Exception: |
| 139 | + logger.exception("Failed to get Mengram cognitive profile.") |
| 140 | + raise |
| 141 | + |
| 142 | + return Data(data=profile) |
0 commit comments