memu-py 1.3.0__cp313-abi3-win_amd64.whl → 1.4.0__cp313-abi3-win_amd64.whl

memu/app/crud.py

@@ -654,7 +654,7 @@ class CRUDMixin:
             prompt = self._build_category_patch_prompt(
                 category=cat, content_before=content_before, content_after=content_after
             )
-            tasks.append(client.summarize(prompt, system_prompt=None))
+            tasks.append(client.chat(prompt))
             target_ids.append(cid)
         if not tasks:
             return

memu/app/memorize.py

@@ -240,7 +240,7 @@ class MemorizeMixin:
         resources: list[Resource] = []
         items: list[MemoryItem] = []
         relations: list[CategoryItem] = []
-        category_updates: dict[str, list[str]] = {}
+        category_updates: dict[str, list[tuple[str, str]]] = {}
         user_scope = state.get("user", {})
 
         for plan in state.get("resource_plans", []):
@@ -282,12 +282,18 @@ class MemorizeMixin:
 
     async def _memorize_persist_and_index(self, state: WorkflowState, step_context: Any) -> WorkflowState:
         llm_client = self._get_step_llm_client(step_context)
-        await self._update_category_summaries(
+        updated_summaries = await self._update_category_summaries(
             state.get("category_updates", {}),
             ctx=state["ctx"],
             store=state["store"],
             llm_client=llm_client,
         )
+        if self.memorize_config.enable_item_references:
+            await self._persist_item_references(
+                updated_summaries=updated_summaries,
+                category_updates=state.get("category_updates", {}),
+                store=state["store"],
+            )
         return state
 
     def _memorize_build_response(self, state: WorkflowState, step_context: Any) -> WorkflowState:
@@ -522,7 +528,8 @@ class MemorizeMixin:
             for mtype in memory_types
         ]
         valid_prompts = [prompt for prompt in prompts if prompt.strip()]
-        tasks = [client.summarize(prompt_text) for prompt_text in valid_prompts]
+        # These prompts are instructions that request structured output, not text summaries.
+        tasks = [client.chat(prompt_text) for prompt_text in valid_prompts]
         responses = await asyncio.gather(*tasks)
         return self._parse_structured_entries(memory_types, responses)
 
@@ -577,14 +584,23 @@ class MemorizeMixin:
         store: Database,
         embed_client: Any | None = None,
         user: Mapping[str, Any] | None = None,
-    ) -> tuple[list[MemoryItem], list[CategoryItem], dict[str, list[str]]]:
+    ) -> tuple[list[MemoryItem], list[CategoryItem], dict[str, list[tuple[str, str]]]]:
+        """
+        Persist memory items and track category updates.
+
+        Returns:
+            Tuple of (items, relations, category_updates)
+            where category_updates maps category_id -> list of (item_id, summary) tuples
+        """
         summary_payloads = [content for _, content, _ in structured_entries]
         client = embed_client or self._get_llm_client()
         item_embeddings = await client.embed(summary_payloads) if summary_payloads else []
         items: list[MemoryItem] = []
         rels: list[CategoryItem] = []
-        category_memory_updates: dict[str, list[str]] = {}
+        # Changed: now stores (item_id, summary) tuples for reference support
+        category_memory_updates: dict[str, list[tuple[str, str]]] = {}
 
+        reinforce = self.memorize_config.enable_item_reinforcement
         for (memory_type, summary_text, cat_names), emb in zip(structured_entries, item_embeddings, strict=True):
             item = store.memory_item_repo.create_item(
                 resource_id=resource_id,
@@ -592,12 +608,17 @@ class MemorizeMixin:
                 summary=summary_text,
                 embedding=emb,
                 user_data=dict(user or {}),
+                reinforce=reinforce,
             )
             items.append(item)
+            if reinforce and item.extra.get("reinforcement_count", 1) > 1:
+                # existing item
+                continue
             mapped_cat_ids = self._map_category_names_to_ids(cat_names, ctx)
             for cid in mapped_cat_ids:
                 rels.append(store.category_item_repo.link_item_category(item.id, cid, user_data=dict(user or {})))
-                category_memory_updates.setdefault(cid, []).append(summary_text)
+                # Store (item_id, summary) tuple for reference support
+                category_memory_updates.setdefault(cid, []).append((item.id, summary_text))
 
         return items, rels, category_memory_updates
 
@@ -779,7 +800,7 @@ class MemorizeMixin:
         preprocessed_text = format_conversation_for_preprocess(text)
         prompt = template.format(conversation=self._escape_prompt_value(preprocessed_text))
         client = llm_client or self._get_llm_client()
-        processed = await client.summarize(prompt, system_prompt=None)
+        processed = await client.chat(prompt)
         _conv, segments = self._parse_conversation_preprocess_with_segments(processed, preprocessed_text)
 
         # Important: always use the original JSON-derived, indexed conversation text for downstream
@@ -809,16 +830,13 @@ class MemorizeMixin:
 
     async def _summarize_segment(self, segment_text: str, llm_client: Any | None = None) -> str | None:
         """Summarize a single conversation segment."""
-        prompt = f"""Summarize the following conversation segment in 1-2 concise sentences.
-Focus on the main topic or theme discussed.
-
-Conversation:
-{segment_text}
-
-Summary:"""
+        system_prompt = (
+            "Summarize the given conversation segment in 1-2 concise sentences. "
+            "Focus on the main topic or theme discussed."
+        )
         try:
             client = llm_client or self._get_llm_client()
-            response = await client.summarize(prompt, system_prompt=None)
+            response = await client.chat(segment_text, system_prompt=system_prompt)
             return response.strip() if response else None
         except Exception:
             logger.exception("Failed to summarize segment")
@@ -895,7 +913,7 @@ Summary:"""
         """Preprocess document data - condense and extract caption"""
         prompt = template.format(document_text=self._escape_prompt_value(text))
         client = llm_client or self._get_llm_client()
-        processed = await client.summarize(prompt, system_prompt=None)
+        processed = await client.chat(prompt)
         processed_content, caption = self._parse_multimodal_response(processed, "processed_content", "caption")
         return [{"text": processed_content or text, "caption": caption}]
 
@@ -905,7 +923,7 @@ Summary:"""
         """Preprocess audio data - format transcription and extract caption"""
         prompt = template.format(transcription=self._escape_prompt_value(text))
         client = llm_client or self._get_llm_client()
-        processed = await client.summarize(prompt, system_prompt=None)
+        processed = await client.chat(prompt)
         processed_content, caption = self._parse_multimodal_response(processed, "processed_content", "caption")
         return [{"text": processed_content or text, "caption": caption}]
 
@@ -960,19 +978,115 @@ Summary:"""
         safe_categories = self._escape_prompt_value(categories_str)
         return template.format(resource=safe_resource, categories_str=safe_categories)
 
-    def _build_category_summary_prompt(self, *, category: MemoryCategory, new_memories: list[str]) -> str:
-        new_items_text = "\n".join(f"- {m}" for m in new_memories if m.strip())
+    def _build_item_ref_id(self, item_id: str) -> str:
+        return item_id.replace("-", "")[:6]
+
+    def _extract_refs_from_summaries(self, summaries: dict[str, str]) -> set[str]:
+        """
+        Extract all [ref:xxx] references from summary texts.
+
+        Args:
+            summaries: dict mapping category_id -> summary text
+
+        Returns:
+            Set of all referenced short IDs (the xxx part from [ref:xxx])
+        """
+        from memu.utils.references import extract_references
+
+        refs: set[str] = set()
+        for summary in summaries.values():
+            refs.update(extract_references(summary))
+        return refs
+
+    async def _persist_item_references(
+        self,
+        *,
+        updated_summaries: dict[str, str],
+        category_updates: dict[str, list[tuple[str, str]]],
+        store: Database,
+    ) -> None:
+        """
+        Persist ref_id to items that are referenced in category summaries.
+
+        This function:
+        1. Extracts all [ref:xxx] patterns from updated summaries
+        2. Builds a mapping of short_id -> full item_id for all items in category_updates
+        3. For items whose short_id appears in the references, updates their extra column
+           with {"ref_id": short_id}
+        """
+        # Extract all referenced short IDs from summaries
+        referenced_short_ids = self._extract_refs_from_summaries(updated_summaries)
+        if not referenced_short_ids:
+            return
+
+        # Build mapping of short_id -> full item_id for all items in category_updates
+        short_id_to_item_id: dict[str, str] = {}
+        for item_tuples in category_updates.values():
+            for item_id, _ in item_tuples:
+                short_id = self._build_item_ref_id(item_id)
+                short_id_to_item_id[short_id] = item_id
+
+        # Update extra column for referenced items
+        for short_id in referenced_short_ids:
+            matched_item_id = short_id_to_item_id.get(short_id)
+            if matched_item_id:
+                store.memory_item_repo.update_item(
+                    item_id=matched_item_id,
+                    extra={"ref_id": short_id},
+                )
+
+    def _build_category_summary_prompt(
+        self,
+        *,
+        category: MemoryCategory,
+        new_memories: list[str] | list[tuple[str, str]],
+    ) -> str:
+        """
+        Build the prompt for updating a category summary.
+
+        Args:
+            category: The category to update
+            new_memories: Either list of summary strings (legacy) or list of (item_id, summary) tuples (with refs)
+        """
+        # Check if references are enabled and we have (id, summary) tuples
+        enable_refs = getattr(self.memorize_config, "enable_item_references", False)
+
+        if enable_refs:
+            from memu.prompts.category_summary import (
+                CUSTOM_PROMPT_WITH_REFS as category_summary_custom_prompt,
+            )
+            from memu.prompts.category_summary import (
+                PROMPT_WITH_REFS as category_summary_prompt,
+            )
+
+            tuple_memories = cast(list[tuple[str, str]], new_memories)
+            new_items_text = "\n".join(
+                f"- [{self._build_item_ref_id(item_id)}] {summary}"
+                for item_id, summary in tuple_memories
+                if summary.strip()
+            )
+        else:
+            category_summary_prompt = CATEGORY_SUMMARY_PROMPT
+            category_summary_custom_prompt = CATEGORY_SUMMARY_CUSTOM_PROMPT
+
+            if new_memories and isinstance(new_memories[0], tuple):
+                tuple_memories = cast(list[tuple[str, str]], new_memories)
+                new_items_text = "\n".join(f"- {summary}" for item_id, summary in tuple_memories if summary.strip())
+            else:
+                str_memories = cast(list[str], new_memories)
+                new_items_text = "\n".join(f"- {m}" for m in str_memories if m.strip())
+
         original = category.summary or ""
         category_config = self.category_config_map.get(category.name)
         configured_prompt = (
             category_config and category_config.summary_prompt
         ) or self.memorize_config.default_category_summary_prompt
         if configured_prompt is None:
-            prompt = CATEGORY_SUMMARY_PROMPT
+            prompt = category_summary_prompt
         elif isinstance(configured_prompt, str):
             prompt = configured_prompt
         else:
-            prompt = self._resolve_custom_prompt(configured_prompt, CATEGORY_SUMMARY_CUSTOM_PROMPT)
+            prompt = self._resolve_custom_prompt(configured_prompt, category_summary_custom_prompt)
         target_length = (
             category_config and category_config.target_length
         ) or self.memorize_config.default_category_summary_target_length
@@ -985,13 +1099,20 @@ Summary:"""
 
     async def _update_category_summaries(
         self,
-        updates: dict[str, list[str]],
+        updates: dict[str, list[tuple[str, str]]] | dict[str, list[str]],
         ctx: Context,
         store: Database,
         llm_client: Any | None = None,
-    ) -> None:
+    ) -> dict[str, str]:
+        """
+        Update category summaries based on new memory items.
+
+        Returns:
+            dict mapping category_id -> updated summary text
+        """
+        updated_summaries: dict[str, str] = {}
         if not updates:
-            return
+            return updated_summaries
         tasks = []
         target_ids: list[str] = []
         client = llm_client or self._get_llm_client()
@@ -1000,19 +1121,22 @@ Summary:"""
             if not cat or not memories:
                 continue
             prompt = self._build_category_summary_prompt(category=cat, new_memories=memories)
-            tasks.append(client.summarize(prompt, system_prompt=None))
+            tasks.append(client.chat(prompt))
             target_ids.append(cid)
         if not tasks:
-            return
+            return updated_summaries
         summaries = await asyncio.gather(*tasks)
         for cid, summary in zip(target_ids, summaries, strict=True):
             cat = store.memory_category_repo.categories.get(cid)
             if not cat:
                 continue
+            cleaned_summary = summary.replace("```markdown", "").replace("```", "").strip()
             store.memory_category_repo.update_category(
                 category_id=cid,
-                summary=summary.replace("```markdown", "").replace("```", "").strip(),
+                summary=cleaned_summary,
             )
+            updated_summaries[cid] = cleaned_summary
+        return updated_summaries
 
     def _parse_conversation_preprocess(self, raw: str) -> tuple[str | None, str | None]:
         conversation = self._extract_tag_content(raw, "conversation")

memu/app/patch.py

@@ -407,7 +407,7 @@ class PatchMixin:
             prompt = self._build_category_patch_prompt(
                 category=cat, content_before=content_before, content_after=content_after
             )
-            tasks.append(client.summarize(prompt, system_prompt=None))
+            tasks.append(client.chat(prompt))
             target_ids.append(cid)
         if not tasks:
             return

memu/app/retrieve.py

@@ -321,6 +321,28 @@ class RetrieveMixin:
             state["query_vector"] = (await embed_client.embed([state["active_query"]]))[0]
         return state
 
+    def _extract_referenced_item_ids(self, state: WorkflowState) -> set[str]:
+        """Extract item IDs from category summary references."""
+        from memu.utils.references import extract_references
+
+        category_hits = state.get("category_hits") or []
+        summary_lookup = state.get("category_summary_lookup", {})
+        category_pool = state.get("category_pool") or {}
+        referenced_item_ids: set[str] = set()
+
+        for cid, _score in category_hits:
+            # Get summary from lookup or category
+            summary = summary_lookup.get(cid)
+            if not summary:
+                cat = category_pool.get(cid)
+                if cat:
+                    summary = cat.summary
+            if summary:
+                refs = extract_references(summary)
+                referenced_item_ids.update(refs)
+
+        return referenced_item_ids
+
     async def _rag_recall_items(self, state: WorkflowState, step_context: Any) -> WorkflowState:
         if not state.get("retrieve_item") or not state.get("needs_retrieval") or not state.get("proceed_to_items"):
             state["item_hits"] = []
@@ -338,6 +360,8 @@ class RetrieveMixin:
             qvec,
             self.retrieve_config.item.top_k,
             where=where_filters,
+            ranking=self.retrieve_config.item.ranking,
+            recency_decay_days=self.retrieve_config.item.recency_decay_days,
         )
         state["item_pool"] = items_pool
         return state
@@ -594,10 +618,26 @@ class RetrieveMixin:
             return state
 
         where_filters = state.get("where") or {}
-        category_ids = [cat["id"] for cat in state.get("category_hits", [])]
+        category_hits = state.get("category_hits", [])
+        category_ids = [cat["id"] for cat in category_hits]
         llm_client = self._get_step_llm_client(step_context)
         store = state["store"]
-        items_pool = store.memory_item_repo.list_items(where_filters)
+
+        use_refs = getattr(self.retrieve_config.item, "use_category_references", False)
+        ref_ids: list[str] = []
+        if use_refs and category_hits:
+            # Extract all ref_ids from category summaries
+            from memu.utils.references import extract_references
+
+            for cat in category_hits:
+                summary = cat.get("summary") or ""
+                ref_ids.extend(extract_references(summary))
+        if ref_ids:
+            # Query items by ref_ids
+            items_pool = store.memory_item_repo.list_items_by_ref_ids(ref_ids, where_filters)
+        else:
+            items_pool = store.memory_item_repo.list_items(where_filters)
+
         relations = store.category_item_repo.list_relations(where_filters)
         category_pool = state.get("category_pool") or store.memory_category_repo.list_categories(where_filters)
         state["item_hits"] = await self._llm_rank_items(
@@ -737,7 +777,7 @@ class RetrieveMixin:
 
         sys_prompt = system_prompt or PRE_RETRIEVAL_SYSTEM_PROMPT
         client = llm_client or self._get_llm_client()
-        response = await client.summarize(user_prompt, system_prompt=sys_prompt)
+        response = await client.chat(user_prompt, system_prompt=sys_prompt)
         decision = self._extract_decision(response)
         rewritten = self._extract_rewritten_query(response) or query
 
@@ -1195,7 +1235,7 @@ class RetrieveMixin:
         )
 
         client = llm_client or self._get_llm_client()
-        llm_response = await client.summarize(prompt, system_prompt=None)
+        llm_response = await client.chat(prompt)
         return self._parse_llm_category_response(llm_response, store, categories=category_pool)
 
     async def _llm_rank_items(
@@ -1234,7 +1274,7 @@ class RetrieveMixin:
         )
 
         client = llm_client or self._get_llm_client()
-        llm_response = await client.summarize(prompt, system_prompt=None)
+        llm_response = await client.chat(prompt)
         return self._parse_llm_item_response(llm_response, store, items=item_pool)
 
     async def _llm_rank_resources(
@@ -1279,7 +1319,7 @@ class RetrieveMixin:
         )
 
         client = llm_client or self._get_llm_client()
-        llm_response = await client.summarize(prompt, system_prompt=None)
+        llm_response = await client.chat(prompt)
         return self._parse_llm_resource_response(llm_response, store, resources=resource_pool)
 
     def _parse_llm_category_response(

memu/app/settings.py

@@ -151,6 +151,20 @@ class RetrieveCategoryConfig(BaseModel):
 class RetrieveItemConfig(BaseModel):
     enabled: bool = Field(default=True, description="Whether to enable item retrieval.")
     top_k: int = Field(default=5, description="Total number of items to retrieve.")
+    # Reference-aware retrieval
+    use_category_references: bool = Field(
+        default=False,
+        description="When category retrieval is insufficient, follow [ref:ITEM_ID] citations to fetch referenced items.",
+    )
+    # Salience-aware retrieval settings
+    ranking: Literal["similarity", "salience"] = Field(
+        default="similarity",
+        description="Ranking strategy: 'similarity' (cosine only) or 'salience' (weighted by reinforcement + recency).",
+    )
+    recency_decay_days: float = Field(
+        default=30.0,
+        description="Half-life in days for recency decay in salience scoring. After this many days, recency factor is ~0.5.",
+    )
 
 
 class RetrieveResourceConfig(BaseModel):
@@ -217,6 +231,15 @@ class MemorizeConfig(BaseModel):
         description="Target max length for auto-generated category summaries.",
     )
     category_update_llm_profile: str = Field(default="default", description="LLM profile for category summary.")
+    # Reference tracking for category summaries
+    enable_item_references: bool = Field(
+        default=False,
+        description="Enable inline [ref:ITEM_ID] citations in category summaries linking to source memory items.",
+    )
+    enable_item_reinforcement: bool = Field(
+        default=False,
+        description="Enable reinforcement tracking for memory items.",
+    )
 
 
 class PatchConfig(BaseModel):
@@ -225,6 +248,9 @@ class PatchConfig(BaseModel):
 
 class DefaultUserModel(BaseModel):
     user_id: str | None = None
+    # Agent/session scoping for multi-agent and multi-session memory filtering
+    # agent_id: str | None = None
+    # session_id: str | None = None
 
 
 class UserConfig(BaseModel):

memu/client/__init__.py

@@ -0,0 +1,26 @@
+"""
+memU Client Wrapper for Auto-Recall Memory Injection.
+
+This module provides optional wrappers around OpenAI/Anthropic clients
+that automatically inject recalled memories into prompts.
+
+Usage:
+    from memu.client import wrap_openai
+    from openai import OpenAI
+
+    client = OpenAI()
+    service = MemoryService(...)
+
+    # Wrap the client for auto-recall
+    wrapped_client = wrap_openai(client, service, user_id="user123")
+
+    # Now all chat completions automatically include relevant memories
+    response = wrapped_client.chat.completions.create(
+        model="gpt-4",
+        messages=[{"role": "user", "content": "What's my favorite drink?"}]
+    )
+"""
+
+from memu.client.openai_wrapper import MemuOpenAIWrapper, wrap_openai
+
+__all__ = ["MemuOpenAIWrapper", "wrap_openai"]