Method query

methods/amem_agent.py:302–360

Query the A-Mem memory system and generate response. Steps: 1. Retrieve related memories using A-Mem's find_related_memories 2. Construct context with retrieved memories 3. Generate response using LLM

(
        self,
        question: str,
        system_message: Optional[str] = None,
        **kwargs,
    )

        )

    def query(
        self,
        question: str,
        system_message: Optional[str] = None,
        **kwargs,
    ) -> AgentResponse:
        """Query the A-Mem memory system and generate response.

        Steps:
        1. Retrieve related memories using A-Mem's find_related_memories
        2. Construct context with retrieved memories
        3. Generate response using LLM
        """
        context_id = self._get_context_id()
        memory_system = self._get_memory_system(context_id)

        # Retrieve related memories
        memory_str, indices = memory_system.find_related_memories(question, k=self.retrieve_num)

        # Calculate available tokens for memory context
        system_tokens = self._llm_client.count_tokens(system_message) if system_message else 0
        question_tokens = self._llm_client.count_tokens(question)
        reserved_tokens = self.amem_max_tokens + RESERVED_OUTPUT_TOKENS
        max_memory_tokens = max(
            self.amem_max_context_tokens - system_tokens - question_tokens - reserved_tokens,
            0,
        )

        # Truncate memory if needed and construct full question
        if memory_str.strip():
            memory_str = self._truncate_to_token_limit(memory_str, max_memory_tokens)
            memory_context = f"[Retrieved A-Mem Notes]\n{memory_str}\n\n"
            full_question = memory_context + question
        else:
            full_question = question

        # Generate response
        messages = format_messages(full_question, system_message)
        response = self._llm_client.chat(messages)

        # Build retrieved memories list for logging
        indices_list = indices.tolist() if hasattr(indices, "tolist") else list(indices)
        retrieved_memories: List[Dict[str, Any]] = []
        if memory_str.strip():
            retrieved_memories.append({
                "memory": memory_str[:2000],
                "type": "amem_retrieval",
                "indices": indices_list,
            })

        return AgentResponse(
            output=response.content,
            retrieved_count=len(indices_list),
            retrieved_memories=retrieved_memories,
            extra={
                "method": "amem",
                "context_id": context_id,
            },