Enhance agent and models for improved JSON handling and memory type normalization

- Updated `agent.py` to include a new function for parsing LLM output, ensuring strict JSON compliance and handling various output formats.
- Modified response format in `agent.py` to clarify JSON structure requirements.
- Enhanced `models.py` with a normalization function for memory types, allowing for flexible input handling and improved validation.
- Introduced a field validator in `NewMemoryItem` to ensure correct memory type coercion during model instantiation.

agent.py

@@ -1,7 +1,11 @@
+import json
 import os
+import re
+from json import JSONDecoder
 from typing import List, Dict
 
 from dotenv import load_dotenv
+from pydantic import ValidationError
 from langchain_core.messages import AIMessage, HumanMessage
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain_openai import ChatOpenAI
@@ -31,7 +35,7 @@ RESPONSE FORMAT — you MUST return valid JSON only, no other text:
   "new_memories": [
     {{
       "content": "What the current customer said, preserved in their voice as closely as possible",
-      "type": "episodic|semantic|state|procedural",
+      "type": "episodic|semantic|state|procedural (English only, not episodico/semantico)",
       "context_tags": ["tag1", "tag2"],
       "summary": "5-10 word summary for display"
     }}
@@ -60,6 +64,8 @@ current-sounding situation (state), or insight about what helped/hurt in support
 logical troubleshooting steps, and reasonable next steps without inventing protocol numbers, \
 discounts, stock levels, or coverage
 - Respond ONLY with the JSON object — no preamble, no markdown fences
+- Previous assistant turns in the chat history are plain-text summaries for context; \
+your current reply must still be ONLY the JSON object, never duplicate the answer outside JSON
 """
 
 OPENROUTER_API_URL = os.getenv("OPENROUTER_API_URL")
@@ -72,6 +78,41 @@ PROMPT = ChatPromptTemplate.from_messages([
 ])
 
 
+def _parse_agent_llm_output(text: str) -> AgentLLMOutput:
+    """Accept strict JSON or model output with prose before/after the JSON object."""
+    text = (text or "").strip()
+    if not text:
+        raise ValueError("Empty LLM output")
+
+    decoder = JSONDecoder()
+    candidates: List[str] = [text]
+    for match in re.finditer(
+        r"```(?:json)?\s*(\{.*?\})\s*```", text, re.DOTALL | re.IGNORECASE
+    ):
+        candidates.append(match.group(1))
+
+    for raw in candidates:
+        try:
+            return AgentLLMOutput.model_validate(json.loads(raw))
+        except (json.JSONDecodeError, ValueError, ValidationError):
+            continue
+
+    start = 0
+    while True:
+        brace = text.find("{", start)
+        if brace == -1:
+            break
+        try:
+            obj, _ = decoder.raw_decode(text, brace)
+            if isinstance(obj, dict) and "response" in obj:
+                return AgentLLMOutput.model_validate(obj)
+        except (json.JSONDecodeError, ValidationError):
+            pass
+        start = brace + 1
+
+    raise ValueError("No valid AgentLLMOutput JSON found in model response")
+
+
 class Agent:
     def __init__(self, memory_store: MemoryStore):
         self.memory_store = memory_store
@@ -81,15 +122,16 @@ class Agent:
             model=MODEL,
             max_tokens=1500,
             timeout=60.0,
-            extra_body={"chat_template_kwargs": {"enable_thinking": False}},
+            extra_body={
+                "chat_template_kwargs": {"enable_thinking": False},
+                "response_format": {"type": "json_object"},
+            },
             default_headers={
                 "HTTP-Referer": "http://localhost:8000",
                 "X-Title": "Agent Memory Phase 1",
             },
         )
-        self.chain = PROMPT | llm.with_structured_output(
-            AgentLLMOutput, method="json_mode"
-        )
+        self.chain = PROMPT | llm
 
     async def chat(
         self,
@@ -115,11 +157,12 @@ class Agent:
             for turn in conversation_history[-6:]
         ]
 
-        parsed: AgentLLMOutput = await self.chain.ainvoke({
+        raw = await self.chain.ainvoke({
             "memory_context": memory_context,
             "input": message,
             "history": history,
         })
+        parsed = _parse_agent_llm_output(raw.content)
 
         for mem_id in parsed.memories_used:
             self.memory_store.update_access(mem_id)

models.py

@@ -1,6 +1,8 @@
-from pydantic import BaseModel
-from typing import List, Optional, Dict, Any
+import unicodedata
 from enum import Enum
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, field_validator
 
 
 class MemoryType(str, Enum):
@@ -10,6 +12,28 @@ class MemoryType(str, Enum):
     PROCEDURAL = "procedural"
 
 
+_MEMORY_TYPE_ALIASES = {
+    "episodico": MemoryType.EPISODIC,
+    "episodic": MemoryType.EPISODIC,
+    "semantico": MemoryType.SEMANTIC,
+    "semantic": MemoryType.SEMANTIC,
+    "estado": MemoryType.STATE,
+    "state": MemoryType.STATE,
+    "procedural": MemoryType.PROCEDURAL,
+    "procedimental": MemoryType.PROCEDURAL,
+}
+
+
+def normalize_memory_type(value: Any) -> MemoryType:
+    if isinstance(value, MemoryType):
+        return value
+    if value is None or (isinstance(value, str) and not value.strip()):
+        return MemoryType.SEMANTIC
+    key = unicodedata.normalize("NFKD", str(value).strip().lower())
+    key = "".join(c for c in key if not unicodedata.combining(c))
+    return _MEMORY_TYPE_ALIASES.get(key, MemoryType.SEMANTIC)
+
+
 class Memory(BaseModel):
     id: str
     content: str
@@ -31,6 +55,11 @@ class NewMemoryItem(BaseModel):
     context_tags: List[str] = []
     summary: str = ""
 
+    @field_validator("type", mode="before")
+    @classmethod
+    def coerce_memory_type(cls, value: Any) -> MemoryType:
+        return normalize_memory_type(value)
+
 
 class AgentLLMOutput(BaseModel):
     response: str