Method query

methods/zep_agent.py:398–443

Query the agent.

(
        self,
        question: str,
        system_message: Optional[str] = None,
        sub_dataset: str = "default",
        query_id: Optional[str] = None,
        **kwargs
    )

        return response.content

    def query(
        self,
        question: str,
        system_message: Optional[str] = None,
        sub_dataset: str = "default",
        query_id: Optional[str] = None,
        **kwargs
    ) -> AgentResponse:
        """Query the agent."""
        # Retrieve relevant memories
        retrieval_results = self._retrieve(question, sub_dataset)

        # Compose context
        retrieved_context = compose_search_context(
            edges=retrieval_results["edges"],
            nodes=retrieval_results["nodes"],
            context_block=retrieval_results["context_block"],
            episodes=retrieval_results["episodes"],
        )

        # Call LLM to generate response
        response_text = self._llm_response(retrieved_context, question)

        # Save retrieval results
        if query_id is not None:
            self._save_retrieval_context(
                query_id=query_id,
                context_id=self._context_id,
                sub_dataset=sub_dataset,
                retrieved_context=retrieved_context,
                response=response_text,
            )

        return AgentResponse(
            output=response_text,
            query_time=0.0,
            retrieved_count=self.retrieve_num,
            retrieved_memories=[
                {"memory": retrieved_context, "type": "zep_retrieval"}
            ],  # Fix: properly set field
            extra={
                "method": "zep",
                "graph_id": self._current_graph_id,
                "thread_id": self._current_thread_id,
            }
        )

    def _save_retrieval_context(
        self,