Update app.py

app.py

@@ -73,47 +73,72 @@ def create_index() -> None:
     )
 
 def save_memory(purpose: str, content: str) -> List[Dict]:
-    """Save processed content to memory with indexing."""
+    """Save processed content to memory with indexing using agent prompts."""
     uid = uuid.uuid4()
     timestamp = datetime.datetime.now().strftime("%Y-%m-%d--%H-%M-%S")
     
-    # Generate structured data
-    prompt = f"{INDEX_PROMPT}\nData to index:\n{content[:5000]}"  # Truncate for API limits
     try:
-        response = generate_response(prompt, model="meta-llama/llama-4-maverick:free")
-        structured_data = json.loads(response)
+        # Use the structured prompt from agent.py
+        prompt = PREFIX.format(
+            timestamp=timestamp,
+            purpose="Compile the provided data into structured memory"
+        ) + COMPRESS_DATA_PROMPT_SMALL.format(
+            direction="Create structured memory entry",
+            knowledge="",
+            history=content[:5000]  # Truncate for API limits
+        )
+        
+        if VERBOSE:
+            print(LOG_PROMPT.format(prompt))
+        
+        response = generate_response(prompt, model="anthropic/claude-2")
+        
+        if VERBOSE:
+            print(LOG_RESPONSE.format(response))
+        
+        # Enhanced JSON extraction
+        json_str = response.split('```json')[1].split('```')[0] if '```json' in response else response
+        structured_data = json.loads(json_str.strip())
+        
+        # Validate required structure
+        if not all(key in structured_data for key in ["keywords", "title", "content"]):
+            raise ValueError("Missing required fields in structured data")
+            
+        # Create memory entry
+        memory_entry = {
+            **structured_data,
+            "file_name": f"{timestamp}--{uid}.json",
+            "timestamp": str(datetime.datetime.now()),
+            "source": "arxiv" if "arxiv.org" in content.lower() else "web"
+        }
+        
+        # Update main memory file
+        main_url = f"{SAVE_DATA_URL}mem-test2/main.json"
+        m = requests.get(main_url)
+        main_data = json.loads(m.text) if m.status_code == 200 else []
+        main_data.append(memory_entry)
+        
+        # Save to HuggingFace
+        main_path = f"tmp-main-{uid}.json"
+        with open(main_path, "w") as f:
+            json.dump(main_data, f)
+        
+        api.upload_file(
+            path_or_fileobj=main_path,
+            path_in_repo="/mem-test2/main.json",
+            repo_id=REPO_NAME,
+            repo_type="dataset",
+        )
+        
+        # Update search index
+        create_index()
+        
+        return [memory_entry]
+        
     except Exception as e:
-        print(f"Memory processing error: {e}")
+        print(f"{LOG_RESPONSE.format('Memory processing failed: ' + str(e))}")
         return []
-    
-    # Save to memory files
-    memory_entry = {
-        **structured_data,
-        "file_name": f"{timestamp}--{uid}.json",
-        "timestamp": str(datetime.datetime.now())
-    }
-    
-    # Update main memory file
-    main_url = f"{SAVE_DATA_URL}mem-test2/main.json"
-    m = requests.get(main_url)
-    main_data = json.loads(m.text) if m.status_code == 200 else []
-    main_data.append(memory_entry)
-    
-    main_path = f"tmp-main-{uid}.json"
-    with open(main_path, "w") as f:
-        json.dump(main_data, f)
-    
-    api.upload_file(
-        path_or_fileobj=main_path,
-        path_in_repo="/mem-test2/main.json",
-        repo_id=REPO_NAME,
-        repo_type="dataset",
-    )
-    
-    # Update index
-    create_index()
-    
-    return [memory_entry]
+
     
 def fetch_url_content(url: str) -> Tuple[bool, str]:
     """Fetch content from a URL and return status and content."""