From 8c67abf7efb4ca94192872f8d6e5f013ce5fe77e Mon Sep 17 00:00:00 2001 From: X-iZhang Date: Fri, 13 Feb 2026 19:34:07 +0000 Subject: [PATCH] feat(memory): add structured extraction schemas for user profiles, research preferences, and experiment conclusions --- EvoScientist/middleware/__init__.py | 2 + EvoScientist/middleware/memory.py | 107 ++++++++++++++-------------- 2 files changed, 57 insertions(+), 52 deletions(-) diff --git a/EvoScientist/middleware/__init__.py b/EvoScientist/middleware/__init__.py index 4a51e24..45c07e2 100644 --- a/EvoScientist/middleware/__init__.py +++ b/EvoScientist/middleware/__init__.py @@ -7,11 +7,13 @@ Re-exports middleware classes and factory functions so that existing from .memory import ( EvoMemoryMiddleware, EvoMemoryState, + ExtractedMemory, create_memory_middleware, ) __all__ = [ "EvoMemoryMiddleware", "EvoMemoryState", + "ExtractedMemory", "create_memory_middleware", ] diff --git a/EvoScientist/middleware/memory.py b/EvoScientist/middleware/memory.py index d86c327..01e1bbf 100644 --- a/EvoScientist/middleware/memory.py +++ b/EvoScientist/middleware/memory.py @@ -33,6 +33,8 @@ from collections.abc import Awaitable, Callable from contextvars import ContextVar from typing import TYPE_CHECKING, Any, Annotated, NotRequired, cast +from pydantic import BaseModel, Field + from langchain.agents.middleware.types import ( AgentMiddleware, AgentState, @@ -60,6 +62,49 @@ class EvoMemoryState(AgentState): evo_memory_content: NotRequired[Annotated[str, PrivateStateAttr]] +# --------------------------------------------------------------------------- +# Structured extraction schemas +# --------------------------------------------------------------------------- + +class UserProfile(BaseModel): + """Extracted user profile information.""" + name: str | None = Field(None, description="User's name") + role: str | None = Field(None, description="User's role (e.g. researcher, student)") + institution: str | None = Field(None, description="User's institution or organization") + language: str | None = Field(None, description="User's preferred language") + + +class ResearchPreferences(BaseModel): + """Extracted research preference information.""" + primary_domain: str | None = Field(None, description="Primary research domain") + sub_fields: str | None = Field(None, description="Research sub-fields") + preferred_frameworks: str | None = Field(None, description="Preferred software frameworks") + preferred_models: str | None = Field(None, description="Preferred AI/ML models") + hardware: str | None = Field(None, description="Available hardware (GPUs, etc.)") + constraints: str | None = Field(None, description="Resource or time constraints") + + +class ExperimentConclusion(BaseModel): + """Extracted experiment conclusion (only when a complete experiment was run).""" + title: str = Field(description="Experiment name") + question: str | None = Field(None, description="Research question") + method: str | None = Field(None, description="Method summary") + key_result: str | None = Field(None, description="Primary metric or outcome") + conclusion: str | None = Field(None, description="One-line conclusion") + artifacts: str | None = Field(None, description="Report path if any") + + +class ExtractedMemory(BaseModel): + """Structured output schema for memory extraction. + + Only fields with genuinely new information should be populated. + """ + user_profile: UserProfile | None = Field(None, description="New user profile information") + research_preferences: ResearchPreferences | None = Field(None, description="New research preferences") + experiment_conclusion: ExperimentConclusion | None = Field(None, description="Completed experiment conclusion") + learned_preferences: list[str] | None = Field(None, description="New preferences or habits observed") + + # --------------------------------------------------------------------------- # Extraction prompt – sent to a (cheap) LLM to pull structured facts # --------------------------------------------------------------------------- @@ -79,42 +124,9 @@ current memory shown below. {conversation} -Return a JSON object with ONLY the keys that have new information to add. -Omit keys where there is nothing new. Use `null` for unknown values. - -```json -{{ - "user_profile": {{ - "name": "string or null", - "role": "string or null", - "institution": "string or null", - "language": "string or null" - }}, - "research_preferences": {{ - "primary_domain": "string or null", - "sub_fields": "string or null", - "preferred_frameworks": "string or null", - "preferred_models": "string or null", - "hardware": "string or null", - "constraints": "string or null" - }}, - "experiment_conclusion": {{ - "title": "string – experiment name", - "question": "string – research question", - "method": "string – method summary", - "key_result": "string – primary metric/outcome", - "conclusion": "string – one-line conclusion", - "artifacts": "string – report path if any" - }}, - "learned_preferences": [ - "string – each new preference or habit observed" - ] -}} -``` - Rules: -- Only return keys with genuinely new information. -- If nothing new was found, return an empty JSON object: `{{}}` +- Only populate fields with genuinely new information. +- Leave fields as null if there is nothing new. - Do NOT repeat information already in . - For experiment_conclusion, only include if a complete experiment was actually run. - Be concise. Each value should be a short phrase, not a paragraph. @@ -552,33 +564,24 @@ class EvoMemoryMiddleware(AgentMiddleware): conversation="\n".join(conv_parts), ) - @staticmethod - def _parse_extraction_response(text: str) -> dict[str, Any]: - """Parse JSON from an LLM extraction response.""" - import json - - json_match = re.search(r"```(?:json)?\s*([\s\S]*?)```", text) - if json_match: - text = json_match.group(1) - return json.loads(text.strip()) - def _extract(self, model: BaseChatModel, memory: str, messages: list[AnyMessage]) -> dict[str, Any]: - """Run LLM extraction on recent messages.""" + """Run LLM extraction on recent messages using structured output.""" prompt = self._build_extraction_prompt(memory, messages) try: - response = model.invoke(prompt) - text = response.content if isinstance(response.content, str) else str(response.content) - return self._parse_extraction_response(text) + structured_model = model.with_structured_output(ExtractedMemory) + result = structured_model.invoke(prompt) + return result.model_dump(exclude_none=True) except Exception as e: # noqa: BLE001 logger.warning("Memory extraction failed: %s", e) return {} async def _aextract(self, model: BaseChatModel, memory: str, messages: list[AnyMessage]) -> dict[str, Any]: + """Async: Run LLM extraction on recent messages using structured output.""" prompt = self._build_extraction_prompt(memory, messages) try: - response = await model.ainvoke(prompt) - text = response.content if isinstance(response.content, str) else str(response.content) - return self._parse_extraction_response(text) + structured_model = model.with_structured_output(ExtractedMemory) + result = await structured_model.ainvoke(prompt) + return result.model_dump(exclude_none=True) except Exception as e: # noqa: BLE001 logger.warning("Memory extraction failed: %s", e) return {}